fixed OpenAI fetch tasks; added nicer imports

alr_envs/examples/examples_dmc.py

@@ -1,5 +1,5 @@
-from alr_envs.dmc.suite.ball_in_cup.ball_in_cup_mp_wrapper import DMCBallInCupMPWrapper
-from alr_envs.utils.make_env_helpers import make_dmp_env, make_env
+import alr_envs
+from alr_envs.dmc.suite.ball_in_cup.mp_wrapper import MPWrapper
 
 
 def example_dmc(env_id="fish-swim", seed=1, iterations=1000, render=True):
@@ -17,13 +17,12 @@ def example_dmc(env_id="fish-swim", seed=1, iterations=1000, render=True):
     Returns:
 
     """
-    env = make_env(env_id, seed)
+    env = alr_envs.make_env(env_id, seed)
     rewards = 0
     obs = env.reset()
     print("observation shape:", env.observation_space.shape)
     print("action shape:", env.action_space.shape)
 
-    # number of samples(multiple environment steps)
     for i in range(iterations):
         ac = env.action_space.sample()
         obs, reward, done, info = env.step(ac)
@@ -63,7 +62,7 @@ def example_custom_dmc_and_mp(seed=1, iterations=1, render=True):
 
     # Replace this wrapper with the custom wrapper for your environment by inheriting from the MPEnvWrapper.
     # You can also add other gym.Wrappers in case they are needed.
-    wrappers = [DMCBallInCupMPWrapper]
+    wrappers = [MPWrapper]
     mp_kwargs = {
         "num_dof": 2,
         "num_basis": 5,
@@ -84,9 +83,9 @@ def example_custom_dmc_and_mp(seed=1, iterations=1, render=True):
         "episode_length": 1000,
         # "frame_skip": 1
     }
-    env = make_dmp_env(base_env, wrappers=wrappers, seed=seed, mp_kwargs=mp_kwargs, **kwargs)
+    env = alr_envs.make_dmp_env(base_env, wrappers=wrappers, seed=seed, mp_kwargs=mp_kwargs, **kwargs)
     # OR for a deterministic ProMP:
-    # env = make_detpmp_env(base_env, wrappers=wrappers, seed=seed, mp_kwargs=mp_args)
+    # env = alr_envs.make_detpmp_env(base_env, wrappers=wrappers, seed=seed, mp_kwargs=mp_args)
 
     # This renders the full MP trajectory
     # It is only required to call render() once in the beginning, which renders every consecutive trajectory.

alr_envs/examples/examples_general.py

@@ -1,11 +1,9 @@
-import warnings
 from collections import defaultdict
 
 import gym
 import numpy as np
 
-from alr_envs.utils.make_env_helpers import make_env, make_env_rank
-from alr_envs.utils.mp_env_async_sampler import AlrContextualMpEnvSampler, AlrMpEnvSampler, DummyDist
+import alr_envs
 
 
 def example_general(env_id="Pendulum-v0", seed=1, iterations=1000, render=True):
@@ -23,7 +21,7 @@ def example_general(env_id="Pendulum-v0", seed=1, iterations=1000, render=True):
 
     """
 
-    env = make_env(env_id, seed)
+    env = alr_envs.make_env(env_id, seed)
     rewards = 0
     obs = env.reset()
     print("Observation shape: ", env.observation_space.shape)
@@ -58,7 +56,7 @@ def example_async(env_id="alr_envs:HoleReacher-v0", n_cpu=4, seed=int('533D', 16
     Returns: Tuple of (obs, reward, done, info) with type np.ndarray
 
     """
-    env = gym.vector.AsyncVectorEnv([make_env_rank(env_id, seed, i) for i in range(n_cpu)])
+    env = gym.vector.AsyncVectorEnv([alr_envs.make_env_rank(env_id, seed, i) for i in range(n_cpu)])
     # OR
     # envs = gym.vector.AsyncVectorEnv([make_env(env_id, seed + i) for i in range(n_cpu)])
 

alr_envs/examples/examples_motion_primitives.py

@@ -1,4 +1,4 @@
-from alr_envs import HoleReacherMPWrapper
+from alr_envs import MPWrapper
 from alr_envs.utils.make_env_helpers import make_dmp_env, make_env
 
 
@@ -113,7 +113,7 @@ def example_fully_custom_mp(seed=1, iterations=1, render=True):
 
     # Replace this wrapper with the custom wrapper for your environment by inheriting from the MPEnvWrapper.
     # You can also add other gym.Wrappers in case they are needed.
-    wrappers = [HoleReacherMPWrapper]
+    wrappers = [MPWrapper]
     mp_kwargs = {
         "num_dof": 5,
         "num_basis": 5,

alr_envs/examples/pd_control_gain_tuning.py

@@ -0,0 +1,74 @@
+import numpy as np
+from matplotlib import pyplot as plt
+
+from alr_envs import dmc
+from alr_envs.utils.make_env_helpers import make_detpmp_env
+
+# This might work for some environments, however, please verify either way the correct trajectory information
+# for your environment are extracted below
+SEED = 10
+env_id = "cartpole-swingup"
+wrappers = [dmc.suite.cartpole.MPWrapper]
+
+mp_kwargs = {
+    "num_dof": 1,
+    "num_basis": 5,
+    "duration": 2,
+    "width": 0.025,
+    "policy_type": "motor",
+    "weights_scale": 0.2,
+    "zero_start": True,
+    "policy_kwargs": {
+        "p_gains": 10,
+        "d_gains": 10 # a good starting point is the sqrt of p_gains
+    }
+}
+
+kwargs = dict(time_limit=2, episode_length=200)
+
+env = make_detpmp_env(env_id, wrappers, seed=SEED, mp_kwargs=mp_kwargs,
+                      **kwargs)
+
+# Plot difference between real trajectory and target MP trajectory
+env.reset()
+pos, vel = env.mp_rollout(env.action_space.sample())
+
+base_shape = env.full_action_space.shape
+actual_pos = np.zeros((len(pos), *base_shape))
+actual_pos_ball = np.zeros((len(pos), *base_shape))
+actual_vel = np.zeros((len(pos), *base_shape))
+act = np.zeros((len(pos), *base_shape))
+
+for t, pos_vel in enumerate(zip(pos, vel)):
+    actions = env.policy.get_action(pos_vel[0], pos_vel[1])
+    actions = np.clip(actions, env.full_action_space.low, env.full_action_space.high)
+    _, _, _, _ = env.env.step(actions)
+    act[t, :] = actions
+    # TODO verify for your environment
+    actual_pos[t, :] = env.current_pos
+    # actual_pos_ball[t, :] = env.physics.data.qpos[2:]
+    actual_vel[t, :] = env.current_vel
+
+plt.figure(figsize=(15, 5))
+
+plt.subplot(131)
+plt.title("Position")
+plt.plot(actual_pos, c='C0', label=["true" if i == 0 else "" for i in range(np.prod(base_shape))])
+# plt.plot(actual_pos_ball, label="true pos ball")
+plt.plot(pos, c='C1', label=["MP" if i == 0 else "" for i in range(np.prod(base_shape))])
+plt.xlabel("Episode steps")
+plt.legend()
+
+plt.subplot(132)
+plt.title("Velocity")
+plt.plot(actual_vel, c='C0', label=[f"true" if i == 0 else "" for i in range(np.prod(base_shape))])
+plt.plot(vel, c='C1', label=[f"MP" if i == 0 else "" for i in range(np.prod(base_shape))])
+plt.xlabel("Episode steps")
+plt.legend()
+
+plt.subplot(133)
+plt.title("Actions")
+plt.plot(act, c="C0"),  # label=[f"actions" if i == 0 else "" for i in range(np.prod(base_action_shape))])
+plt.xlabel("Episode steps")
+# plt.legend()
+plt.show()