fix(collision guard): back compatibility truncate episode

For backward compatibility with tests there is a default true option which truncates the episode when a collision occurred

Author: juelg

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
@@ -80,6 +82,11 @@ def step(self, action: dict[str, Any]) -> tuple[dict[str, Any], SupportsFloat, b
         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
@@ -114,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)
@@ -140,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,
+        )