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)
 7    env.reset(seed=seed)
 8    for i in range(iterations):
 9        done = False
10        while done is False:
11            ac = env.action_space.sample()
12            obs, reward, terminated, truncated, info = env.step(ac)
13            if render:
14                env.render(mode="human")
15            if terminated or truncated:
16                env.reset()
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 = "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    env = fancy_gym.make_bb(env_id=base_env_id, wrappers=wrappers, black_box_kwargs=black_box_kwargs,
42                            traj_gen_kwargs=trajectory_generator_kwargs, controller_kwargs=controller_kwargs,
43                            phase_kwargs=phase_generator_kwargs, basis_kwargs=basis_generator_kwargs,
44                            seed=seed)
45    if render:
46        env.render(mode="human")
47
48    obs = env.reset()
49
50    for i in range(iteration):
51        ac = env.action_space.sample()
52        obs, reward, terminated, truncated, info = env.step(ac)
53        if terminated or truncated:
54            env.reset()
55
56    env.close()
57    del env
58
59
60if __name__ == "__main__":
61    # run a registered replanning environment
62    example_run_replanning_env(env_name="fancy_ProDMP/BoxPushingDenseReplan-v0", seed=1, iterations=1, render=False)
63
64    # run a custom replanning environment
65    example_custom_replanning_envs(seed=0, iteration=8, render=True)