Best way to reset-to-state from original diffusion policy dataset?

[c0g]

If I call env.reset(options={'reset_to_state': state) using a state from pusht_cchi_v7_replay.zarr, the replayed data looks wrong. If I do:

    env.reset()
    env.unwrapped.agent.position = list(state[0, :2])
    env.unwrapped.block.angle = state[0, 4]
    env.unwrapped.block.position = list(state[0, 2:4])
    env.unwrapped.space.step(env.unwrapped.dt)

it looks right. This is similar to a change and I think is flagged in the docs. Is there a blessed/better way to handle resetting to a known starting state?