from osim.env import L2RunEnv
import opensim

env = L2RunEnv(visualize=True)
observation = env.reset()

s = 0
for s in range(80):
    d = False

    if s == 30:
        state_old = env.osim_model.get_state()
        print("State stored")
        print(state_old)
    if s % 50 == 49:
        env.osim_model.set_state(state_old)
        print("Rollback")
        print(state_old)

    o, r, d, i = env.step(env.action_space.sample())