bug fixes

mani_skill2/agents/controllers/pd_ee_pose.py

@@ -28,6 +28,7 @@
 # NOTE(jigu): not necessary to inherit, just for convenience
 class PDEEPosController(PDJointPosController):
     config: "PDEEPosControllerConfig"
+    _target_pose = None
 
     def _initialize_joints(self):
         self.initial_qpos = None

mani_skill2/agents/controllers/pd_joint_pos.py

@@ -13,6 +13,8 @@
 
 class PDJointPosController(BaseController):
     config: "PDJointPosControllerConfig"
+    _start_qpos = None
+    _target_qpos = None
 
     def _get_joint_limits(self):
         qlimits = self.articulation.get_qlimits()[0, self.joint_indices].cpu().numpy()
@@ -44,8 +46,19 @@ def set_drive_property(self):
     def reset(self):
         super().reset()
         self._step = 0  # counter of simulation steps after action is set
-        self._start_qpos[self.scene._reset_mask] = self.qpos[self.scene._reset_mask]
-        self._target_qpos[self.scene._reset_mask] = self.qpos[self.scene._reset_mask]
+        if self._start_qpos is None:
+            self._start_qpos = self.qpos.clone()
+        else:
+
+            self._start_qpos[self.scene._reset_mask] = self.qpos[
+                self.scene._reset_mask
+            ].clone()
+        if self._target_qpos is None:
+            self._target_qpos = self.qpos.clone()
+        else:
+            self._target_qpos[self.scene._reset_mask] = self.qpos[
+                self.scene._reset_mask
+            ].clone()
 
     def set_drive_targets(self, targets):
         self.articulation.set_joint_drive_targets(
@@ -63,8 +76,10 @@ def set_action(self, action: Array):
             else:
                 self._target_qpos = self._start_qpos + action
         else:
-            # Compatible with mimic controllers
-            self._target_qpos = torch.broadcast_to(action, self._start_qpos.shape)
+            # Compatible with mimic controllers. Need to clone here otherwise cannot do in-place replacements in the reset function
+            self._target_qpos = torch.broadcast_to(
+                action, self._start_qpos.shape
+            ).clone()
         if self.config.interpolate:
             self._step_size = (self._target_qpos - self._start_qpos) / self._sim_steps
         else:

mani_skill2/agents/controllers/pd_joint_pos_vel.py

@@ -10,6 +10,7 @@
 
 class PDJointPosVelController(PDJointPosController):
     config: "PDJointPosVelControllerConfig"
+    _target_qvel = None
 
     def _initialize_action_space(self):
         joint_limits = self._get_joint_limits()

manualtest/record_test.py

@@ -10,7 +10,7 @@
 if __name__ == "__main__":
     # sapien.set_log_level("info")
     # , "StackCube-v1", "PickCube-v1", "PushCube-v1", "PickSingleYCB-v1", "OpenCabinet-v1"
-    num_envs = 2
+    num_envs = 16
     for env_id in ["PickCube-v1"]:
         env = gym.make(
             env_id,
@@ -46,10 +46,10 @@
         env.reset()
         # for i in range(180):
         #     env.step(env.action_space.sample())
-        # env.step(env.action_space.sample())
-        # env.step(env.action_space.sample())
+        env.step(env.action_space.sample())
+        env.step(env.action_space.sample())
         # print("partial reset")
-        # env.reset(options=dict(env_idx=[0]))
+        env.reset(options=dict(env_idx=[0, 1, 2]))
         for i in range(50):
             env.step(env.action_space.sample())
         # for i in range(60):

tests/test_wrappers.py

@@ -150,6 +150,7 @@ def test_recordepisode_wrapper_partial_reset_gpu(env_id, obs_mode):
     for i in range(20):
         obs, rew, terminated, truncated, info = env.step(action_space.sample())
         if i == 13:
+            # should observe in videos (which are organized column by column order) 0, 1, 14, 15 get reset in the middle
             env.reset(options=dict(env_idx=[0, 1, 14, 15]))
     env.close()
     del env