How do I reset the simulation state?

[henrycharlesworth]

How can I reset the simulation state back to a given point? In the other Gym robotics environments I would do:

state = env.sim.get_state()

#restore
env.sim.set_state(state)
env.sim.forward()

I've tried something similar here, but I'm not getting the right results (i.e. obs isn't the same as obs_2, r isn't the same as r_2)

import numpy as np
from robogym.envs.dactyl.full_perpendicular import FullPerpendicularEnv

env = FullPerpendicularEnv.build(constants={"goal_generation": "face_cube_solver"}, apply_wrappers=False)
action = np.random.randn(20)

env.reset()
state = env.mujoco_simulation.get_state()

obs, r, d, i = env.step(action)

env.mujoco_simulation.set_state(state)
env.mujoco_simulation.forward()

obs_2, r_2, d_2, i_2 = env.step(action)

I'll look into this further but any ideas what else needs to be done here?

EDIT: so the things that seem to be different upon loading the simulation state are

env.mujoco_simulation.mj_sim.data.actuator_force
env.mujoco_simulation.mj_sim.data.ctrl

Everything else (all the qpos, qvel) seems to be correct. Not sure what's causing this though? But manually changing these also doesn't work, so there must be something else too.