fix(collision guard): hard sim update and identity action

python/rcsss/_core/sim.pyi

@@ -89,6 +89,7 @@ class FR3(rcsss._core.common.Robot):
     def __init__(self, sim: Sim, id: str, ik: rcsss._core.common.IK) -> None: ...
     def get_parameters(self) -> FR3Config: ...
     def get_state(self) -> FR3State: ...
+    def set_joints_hard(self, q: numpy.ndarray[typing.Literal[7], numpy.dtype[numpy.float64]]) -> None: ...
     def set_parameters(self, cfg: FR3Config) -> bool: ...
 
 class FR3Config(rcsss._core.common.RConfig):

python/rcsss/envs/sim.py

@@ -49,6 +49,7 @@ def __init__(
         check_home_collision: bool = True,
         to_joint_control: bool = False,
         sim_gui: bool = True,
+        truncate_on_collision: bool = True,
     ):
         super().__init__(env)
         self.unwrapped: FR3Env
@@ -58,6 +59,7 @@ def __init__(
         self._logger = logging.getLogger(__name__)
         self.check_home_collision = check_home_collision
         self.to_joint_control = to_joint_control
+        self.truncate_on_collision = truncate_on_collision
         if to_joint_control:
             assert (
                 self.unwrapped.get_unwrapped_control_mode(-2) == ControlMode.JOINTS
@@ -68,23 +70,23 @@ def __init__(
             self.sim.open_gui()
 
     def step(self, action: dict[str, Any]) -> tuple[dict[str, Any], SupportsFloat, bool, bool, dict[str, Any]]:
-        # TODO: we should set the state of the sim to the state of the real robot
+
+        self.collision_env.get_wrapper_attr("robot").set_joints_hard(self.unwrapped.robot.get_joint_position())
         _, _, _, _, info = self.collision_env.step(action)
+
         if self.to_joint_control:
             fr3_env = self.collision_env.unwrapped
             assert isinstance(fr3_env, FR3Env), "Collision env must be an FR3Env instance."
             action[self.unwrapped.joints_key] = fr3_env.robot.get_joint_position()
 
-        # modify action to be joint angles down stream
-        if info["collision"] or not info["ik_success"] or not info["is_sim_converged"]:
-            # return old obs, with truncated and print warning
-            self._logger.warning("Collision detected! Truncating episode: %s", info)
-            if self.last_obs is None:
-                msg = "Collisions detected and no old observation."
-                raise RuntimeError(msg)
-            old_obs, old_info = self.last_obs
-            old_info.update(info)
-            return old_obs, 0, False, True, old_info
+        if info["collision"]:
+            self._logger.warning("Collision detected! %s", info)
+            action[self.unwrapped.joints_key] = self.unwrapped.robot.get_joint_position()
+            if self.truncate_on_collision:
+                if self.last_obs is None:
+                    msg = "Collision detected in the first step!"
+                    raise RuntimeError(msg)
+                return self.last_obs[0], 0, True, True, info
 
         obs, reward, done, truncated, info = super().step(action)
         self.last_obs = obs, info
@@ -119,6 +121,7 @@ def env_from_xml_paths(
         tcp_offset: rcsss.common.Pose | None = None,
         control_mode: ControlMode | None = None,
         sim_gui: bool = True,
+        truncate_on_collision: bool = True,
     ) -> "CollisionGuard":
         assert isinstance(env.unwrapped, FR3Env)
         simulation = sim.Sim(mjmld)
@@ -145,4 +148,12 @@ def env_from_xml_paths(
             gripper_cfg = sim.FHConfig()
             fh = sim.FrankaHand(simulation, id, gripper_cfg)
             c_env = GripperWrapper(c_env, fh)
-        return cls(env, simulation, c_env, check_home_collision, to_joint_control, sim_gui)
+        return cls(
+            env=env,
+            simulation=simulation,
+            collision_env=c_env,
+            check_home_collision=check_home_collision,
+            to_joint_control=to_joint_control,
+            sim_gui=sim_gui,
+            truncate_on_collision=truncate_on_collision,
+        )

src/pybind/rcsss.cpp

@@ -470,6 +470,7 @@ PYBIND11_MODULE(_core, m) {
            py::arg("sim"), py::arg("id"), py::arg("ik"))
       .def("get_parameters", &rcs::sim::FR3::get_parameters)
       .def("set_parameters", &rcs::sim::FR3::set_parameters, py::arg("cfg"))
+      .def("set_joints_hard", &rcs::sim::FR3::set_joints_hard, py::arg("q"))
       .def("get_state", &rcs::sim::FR3::get_state);
   py::enum_<rcs::sim::CameraType>(sim, "CameraType")
       .value("free", rcs::sim::CameraType::free)

src/sim/FR3.cpp

@@ -208,6 +208,14 @@ void FR3::m_reset() {
   }
 }
 
+void FR3::set_joints_hard(const common::Vector7d& q) {
+  for (size_t i = 0; i < std::size(this->ids.joints); ++i) {
+    size_t jnt_id = this->ids.joints[i];
+    size_t jnt_qposadr = this->sim->m->jnt_qposadr[jnt_id];
+    this->sim->d->qpos[jnt_qposadr] = q[i];
+  }
+}
+
 common::Pose FR3::get_base_pose_in_world_coordinates() {
   auto id = mj_name2id(this->sim->m, mjOBJ_BODY,
                        (std::string("base_") + this->id).c_str());

src/sim/FR3.h

@@ -50,6 +50,7 @@ class FR3 : public common::Robot {
   common::Pose get_base_pose_in_world_coordinates() override;
   std::optional<std::shared_ptr<common::IK>> get_ik() override;
   void reset() override;
+  void set_joints_hard(const common::Vector7d &q);
 
  private:
   FR3Config cfg;