Fix _apply_action

huggingface · Apr 9, 2024 · ce294c0 · ce294c0
1 parent 08ddd5a
commit ce294c0
Showing 1 changed file with 2 additions and 11 deletions.
diff --git a/gym_xarm/tasks/base.py b/gym_xarm/tasks/base.py
@@ -238,10 +238,8 @@ def _reset_sim(self):
         self.data.time = self.initial_time
         self.data.qpos[:] = np.copy(self.initial_qpos)
         self.data.qvel[:] = np.copy(self.initial_qvel)
-        if self.model.na != 0:
-            self.data.act[:] = None
-
-        mujoco.mj_forward(self.model, self.data)
+        self._sample_goal()
+        self._mujoco.mj_step(self.model, self.data, nstep=10)
         return True
 
     def get_obs(self):
@@ -307,13 +305,6 @@ def _apply_action(self, action):
             np.concatenate([pos_ctrl, self.gripper_rotation, gripper_ctrl]),
         )
 
-        self.data.time = self.initial_time
-        self.data.qpos[:] = np.copy(self.initial_qpos)
-        self.data.qvel[:] = np.copy(self.initial_qvel)
-        self._sample_goal()
-        self._mujoco.mj_step(self.model, self.data, nstep=10)
-        return True
-
     def _set_gripper(self, gripper_pos, gripper_rotation):
         self._utils.set_mocap_pos(self.model, self.data, "robot0:mocap", gripper_pos)
         self._utils.set_mocap_quat(self.model, self.data, "robot0:mocap", gripper_rotation)