1from collections import OrderedDict
2
3import numpy as np
4from matplotlib import pyplot as plt
5
6import gymnasium as gym
7import fancy_gym
8
9# This might work for some environments, however, please verify either way the correct trajectory information
10# for your environment are extracted below
11SEED = 1
12
13env_id = "fancy_ProMP/Reacher5d-v0"
14
15env = fancy_gym.make(env_id, mp_config_override={'controller_kwargs': {'p_gains': 0.05, 'd_gains': 0.05}}).env
16env.action_space.seed(SEED)
17
18# Plot difference between real trajectory and target MP trajectory
19env.reset(seed=SEED)
20w = env.action_space.sample()
21pos, vel = env.get_trajectory(w)
22
23base_shape = env.env.action_space.shape
24actual_pos = np.zeros((len(pos), *base_shape))
25actual_vel = np.zeros((len(pos), *base_shape))
26act = np.zeros((len(pos), *base_shape))
27
28plt.ion()
29fig = plt.figure()
30ax = fig.add_subplot(1, 1, 1)
31
32img = ax.imshow(env.env.render(mode="rgb_array"))
33fig.show()
34
35for t, (des_pos, des_vel) in enumerate(zip(pos, vel)):
36 actions = env.tracking_controller.get_action(des_pos, des_vel, env.current_pos, env.current_vel)
37 actions = np.clip(actions, env.env.action_space.low, env.env.action_space.high)
38 env.env.step(actions)
39 if t % 15 == 0:
40 img.set_data(env.env.render(mode="rgb_array"))
41 fig.canvas.draw()
42 fig.canvas.flush_events()
43 act[t, :] = actions
44 # TODO verify for your environment
45 actual_pos[t, :] = env.current_pos
46 actual_vel[t, :] = env.current_vel
47
48plt.figure(figsize=(15, 5))
49
50plt.subplot(131)
51plt.title("Position")
52p1 = plt.plot(actual_pos, c='C0', label="true")
53p2 = plt.plot(pos, c='C1', label="MP")
54plt.xlabel("Episode steps")
55handles, labels = plt.gca().get_legend_handles_labels()
56
57by_label = OrderedDict(zip(labels, handles))
58plt.legend(by_label.values(), by_label.keys())
59
60plt.subplot(132)
61plt.title("Velocity")
62plt.plot(actual_vel, c='C0', label="true")
63plt.plot(vel, c='C1', label="MP")
64plt.xlabel("Episode steps")
65
66plt.subplot(133)
67plt.title(f"Actions {np.std(act, axis=0)}")
68plt.plot(act, c="C0"),
69plt.xlabel("Episode steps")
70plt.show()