Replanning Example

 1import gymnasium as gym
 2import fancy_gym
 3
 4
 5def example_run_replanning_env(env_name="fancy_ProDMP/BoxPushingDenseReplan-v0", seed=1, iterations=1, render=False):
 6    env = gym.make(env_name, render_mode='human' if render else None)
 7    env.reset(seed=seed)
 8    for i in range(iterations):
 9        while True:
10            ac = env.action_space.sample()
11            obs, reward, terminated, truncated, info = env.step(ac)
12            if render:
13                env.render()
14            if terminated or truncated:
15                env.reset()
16                break
17    env.close()
18    del env
19
20
21def example_custom_replanning_envs(seed=0, iteration=100, render=True):
22    # id for a step-based environment
23    base_env_id = "fancy/BoxPushingDense-v0"
24
25    wrappers = [fancy_gym.envs.mujoco.box_pushing.mp_wrapper.MPWrapper]
26
27    trajectory_generator_kwargs = {'trajectory_generator_type': 'prodmp',
28                                   'weights_scale': 1}
29    phase_generator_kwargs = {'phase_generator_type': 'exp'}
30    controller_kwargs = {'controller_type': 'velocity'}
31    basis_generator_kwargs = {'basis_generator_type': 'prodmp',
32                              'num_basis': 5}
33
34    # max_planning_times: the maximum number of plans can be generated
35    # replanning_schedule: the trigger for replanning
36    # condition_on_desired: use desired state as the boundary condition for the next plan
37    black_box_kwargs = {'max_planning_times': 4,
38                        'replanning_schedule': lambda pos, vel, obs, action, t: t % 25 == 0,
39                        'condition_on_desired': True}
40
41    base_env = gym.make(base_env_id, render_mode='human' if render else None)
42    env = fancy_gym.make_bb(env=base_env, wrappers=wrappers, black_box_kwargs=black_box_kwargs,
43                            traj_gen_kwargs=trajectory_generator_kwargs, controller_kwargs=controller_kwargs,
44                            phase_kwargs=phase_generator_kwargs, basis_kwargs=basis_generator_kwargs,
45                            seed=seed)
46    if render:
47        env.render()
48
49    obs = env.reset()
50
51    for i in range(iteration):
52        ac = env.action_space.sample()
53        obs, reward, terminated, truncated, info = env.step(ac)
54        if terminated or truncated:
55            env.reset()
56
57    env.close()
58    del env
59
60def main(render=False):
61    # run a registered replanning environment
62    example_run_replanning_env(env_name="fancy_ProDMP/BoxPushingDenseReplan-v0", seed=1, iterations=1, render=render)
63
64    # run a custom replanning environment
65    example_custom_replanning_envs(seed=0, iteration=8, render=render)
66
67if __name__ == "__main__":
68    main()