mp_pytorch now running with zero start/goal promp, but delay is not working

alr_envs/alr/mujoco/reacher/alr_reacher.py

@@ -56,19 +56,22 @@ class ALRReacherEnv(MujocoEnv, utils.EzPickle):
         return ob, reward, done, dict(reward_dist=reward_dist, reward_ctrl=reward_ctrl,
                                       velocity=angular_vel, reward_balance=reward_balance,
                                       end_effector=self.get_body_com("fingertip").copy(),
-                                      goal=self.goal if hasattr(self, "goal") else None)
+                                      goal=self.goal if hasattr(self, "goal") else None,
+                                      joint_pos = self.sim.data.qpos.flat[:self.n_links].copy(),
+                                      joint_vel = self.sim.data.qvel.flat[:self.n_links].copy())
 
     def viewer_setup(self):
         self.viewer.cam.trackbodyid = 0
 
     def reset_model(self):
-        qpos = self.np_random.uniform(low=-0.1, high=0.1, size=self.model.nq) + self.init_qpos
+        # qpos = self.np_random.uniform(low=-0.1, high=0.1, size=self.model.nq) + self.init_qpos
+        qpos = self.init_qpos
         while True:
             self.goal = self.np_random.uniform(low=-self.n_links / 10, high=self.n_links / 10, size=2)
             if np.linalg.norm(self.goal) < self.n_links / 10:
                 break
         qpos[-2:] = self.goal
-        qvel = self.init_qvel + self.np_random.uniform(low=-.005, high=.005, size=self.model.nv)
+        qvel = self.init_qvel# + self.np_random.uniform(low=-.005, high=.005, size=self.model.nv)
         qvel[-2:] = 0
         self.set_state(qpos, qvel)
         self._steps = 0

alr_envs/alr/mujoco/reacher/new_mp_wrapper.py

@@ -22,3 +22,6 @@ class MPWrapper(BaseMPWrapper):
             # self.get_body_com("target"),  # only return target to make problem harder
             [False],  # step
         ])
+
+    def _step_callback(self, action):
+        pass