def example():
task = generate_task(task_generator_id='picking')
env = CausalWorld(task=task, enable_visualization=True)
env.set_starting_state(
{'goal_block': {
'cartesian_position': [0.1, 0.1, 0.1]
}})
for _ in range(500):
obs, reward, done, info = env.step(env.action_space.sample())
env.reset_default_state()
for _ in range(500):
obs, reward, done, info = env.step(env.action_space.sample())
env.reset()
for _ in range(500):
obs, reward, done, info = env.step(env.action_space.sample())
env.close()
def test_reset_default_state(self):
task = generate_task(task_generator_id="picking")
env = CausalWorld(task=task, enable_visualization=False, seed=0)
actions = [env.action_space.sample() for _ in range(200)]
observations_1 = []
rewards_1 = []
env.reset()
for i in range(200):
observations, rewards, _, _ = env.step(actions[i])
observations_1.append(observations)
rewards_1.append(rewards)
env.set_starting_state(
{'goal_block': {
'cylindrical_position': [0.1, np.pi, 0.1]
}})
for i in range(200):
observations, rewards, _, _ = env.step(actions[i])
env.reset_default_state()
for i in range(200):
observations, rewards, _, _ = env.step(actions[i])
assert np.array_equal(observations_1[i], observations)
env.close()
return