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

- state of the real robot is setted to the sim in a hard way
- support to set hard joint values in sim
- if collision is detected the current robot state is set as action and
  not no action

Author: juelg

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

@@ -68,23 +68,18 @@ 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()
 
         obs, reward, done, truncated, info = super().step(action)
         self.last_obs = obs, info

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;