fix(sim): CoverWrapper wrapper for correct resetting

There is now a CoverWrapper which needs to be added at the top of the stack. It ensures that the sim is correctly reset before the wrappers assume a clean sim state. Its only strictly necessary for simulator environments but also works for hardware environments.

Author: juelg

README.md

@@ -52,6 +52,7 @@ from rcs.envs.base import (
     RelativeTo,
     RobotWrapper,
     SimEnv,
+    CoverWrapper,
 )
 from rcs.envs.sim import GripperWrapperSim, RobotSimWrapper
 from rcs.envs.utils import (
@@ -98,6 +99,7 @@ if __name__ == "__main__":
 
     # relative actions bounded by 10cm translation and 10 degree rotation
     env = RelativeActionSpace(env, max_mov=(0.1, np.deg2rad(10)), relative_to=RelativeTo.LAST_STEP)
+    env = CoverWrapper(env)
 
     env.get_wrapper_attr("sim").open_gui()
     # wait for gui to open

examples/fr3/fr3_readme.py

@@ -11,6 +11,7 @@
     RelativeTo,
     RobotWrapper,
     SimEnv,
+    CoverWrapper,
 )
 from rcs.envs.sim import GripperWrapperSim, RobotSimWrapper
 from rcs.envs.utils import (
@@ -57,6 +58,7 @@
 
     # relative actions bounded by 10cm translation and 10 degree rotation
     env = RelativeActionSpace(env, max_mov=(0.1, np.deg2rad(10)), relative_to=RelativeTo.LAST_STEP)
+    env = CoverWrapper(env)
 
     env.get_wrapper_attr("sim").open_gui()
     # wait for gui to open

extensions/rcs_fr3/src/rcs_fr3/creators.py

@@ -18,6 +18,7 @@
     RelativeActionSpace,
     RelativeTo,
     RobotWrapper,
+    CoverWrapper,
 )
 from rcs.envs.creators import RCSHardwareEnvCreator
 from rcs.hand.tilburg_hand import TilburgHand
@@ -125,6 +126,7 @@ def __call__(  # type: ignore
         #     )
         if relative_to != RelativeTo.NONE:
             env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
+        env = CoverWrapper(env)
 
         return env
 
@@ -173,6 +175,7 @@ def __call__(  # type: ignore
             camera_set.wait_for_frames()
             logger.info("CameraSet started")
             env = CameraSetWrapper(env, camera_set)
+        env = CoverWrapper(env)
         return env
 
 

extensions/rcs_panda/src/rcs_panda/creators.py

@@ -14,6 +14,7 @@
     RelativeActionSpace,
     RelativeTo,
     RobotWrapper,
+    CoverWrapper,
 )
 from rcs.envs.creators import RCSHardwareEnvCreator
 from rcs.hand.tilburg_hand import TilburgHand
@@ -106,6 +107,7 @@ def __call__(  # type: ignore
         #     )
         if max_relative_movement is not None:
             env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
+        env = CoverWrapper(env)
 
         return env
 
@@ -153,4 +155,5 @@ def __call__(  # type: ignore
             camera_set.wait_for_frames()
             logger.info("CameraSet started")
             env = CameraSetWrapper(env, camera_set)
+        env = CoverWrapper(env)
         return env

extensions/rcs_so101/src/rcs_so101/creators.py

@@ -10,6 +10,7 @@
     RelativeActionSpace,
     RelativeTo,
     RobotWrapper,
+    CoverWrapper,
 )
 from rcs.envs.creators import RCSHardwareEnvCreator
 from rcs_so101 import SO101IK
@@ -48,6 +49,7 @@ def __call__(  # type: ignore
 
         if max_relative_movement is not None:
             env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
+        env = CoverWrapper(env)
 
         return env
 

extensions/rcs_ur5e/src/rcs_ur5e/creators.py

@@ -10,6 +10,7 @@
     RelativeActionSpace,
     RelativeTo,
     RobotWrapper,
+    CoverWrapper,
 )
 from rcs.envs.creators import RCSHardwareEnvCreator
 from rcs_ur5e.hw import RobotiQGripper, UR5e, UR5eConfig
@@ -51,5 +52,6 @@ def __call__(  # type: ignore
 
         if max_relative_movement is not None:
             env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
+        env = CoverWrapper(env)
 
         return env

extensions/rcs_xarm7/src/rcs_xarm7/creators.py

@@ -12,6 +12,7 @@
     RelativeActionSpace,
     RelativeTo,
     RobotWrapper,
+    CoverWrapper,
 )
 from rcs.envs.creators import RCSHardwareEnvCreator
 from rcs.hand.tilburg_hand import THConfig, TilburgHand
@@ -49,5 +50,6 @@ def __call__(  # type: ignore
 
         if max_relative_movement is not None:
             env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
+        env = CoverWrapper(env)
 
         return env

python/rcs/envs/base.py

@@ -244,8 +244,7 @@ def step_sim(self):
         else:
             self.sim.step_until_convergence()
 
-    def reset_sim(self):
-        self.sim.reset()
+    def apply_sim_state(self):
         self.sim.step(1)
 
 
@@ -255,7 +254,20 @@ def reset(
         if self.main_greenlet is not None:
             self.main_greenlet.switch()
         else:
-            self.reset_sim()
+            self.apply_sim_state()
+        return super().reset(seed=seed, options=options)
+
+class CoverWrapper(gym.Wrapper):
+    """The CoverWrapper must be the last wrapper on the stack
+    
+    Only strictly necessary for simulator environments, but also works for hardware environments.
+    It takes care of resetting the simulator before any other wrapper resets its state, already assuming
+    a fresh simulator state.
+    """
+    def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[dict[str, Any], dict[str, Any]]:
+        if self.env.get_wrapper_attr("PLATFORM") == RobotPlatform.SIMULATION:
+            sim = cast(simulation.Sim, self.get_wrapper_attr("sim"))
+            sim.reset()
         return super().reset(seed=seed, options=options)
 
 
@@ -333,25 +345,28 @@ def action(self, action: dict[str, Any]) -> dict[str, Any]:
         ):
             msg = "Given type is not matching control mode!"
             raise RuntimeError(msg)
+        last_action = self.prev_action
         self.prev_action = copy.deepcopy(action)
 
+        # shallow copy
+        action = dict(action)
         if self.get_base_control_mode() == ControlMode.JOINTS and (
-            self.prev_action is None
-            or not np.allclose(action[self.joints_key], self.prev_action[self.joints_key], atol=1e-03, rtol=0)
+            last_action is None
+            or not np.allclose(action[self.joints_key], last_action[self.joints_key], atol=1e-03, rtol=0)
         ):
             self.robot.set_joint_position(action[self.joints_key])
             action.pop(self.joints_key)
         elif self.get_base_control_mode() == ControlMode.CARTESIAN_TRPY and (
-            self.prev_action is None
-            or not np.allclose(action[self.trpy_key], self.prev_action[self.trpy_key], atol=1e-03, rtol=0)
+            last_action is None
+            or not np.allclose(action[self.trpy_key], last_action[self.trpy_key], atol=1e-03, rtol=0)
         ):
             self.robot.set_cartesian_position(
                 common.Pose(translation=action[self.trpy_key][:3], rpy_vector=action[self.trpy_key][3:])
             )
             action.pop(self.trpy_key)
         elif self.get_base_control_mode() == ControlMode.CARTESIAN_TQuat and (
-            self.prev_action is None
-            or not np.allclose(action[self.tquat_key], self.prev_action[self.tquat_key], atol=1e-03, rtol=0)
+            last_action is None
+            or not np.allclose(action[self.tquat_key], last_action[self.tquat_key], atol=1e-03, rtol=0)
         ):
             self.robot.set_cartesian_position(
                 common.Pose(translation=action[self.tquat_key][:3], quaternion=action[self.tquat_key][3:])
@@ -361,18 +376,13 @@ def action(self, action: dict[str, Any]) -> dict[str, Any]:
 
     def observation(self, observation: dict, info: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
         observation.update(self.get_robot_obs())
-        # if self.env.get_wrapper_attr("PLATFORM") == RobotPlatform.SIMULATION:
-        #     sim_robot = cast(SimRobot, self.robot)
-        #     state = sim_robot.get_state()
-        #     info["collision"] = state.collision
-        #     info["ik_success"] = state.ik_success
-        #     info["is_sim_converged"] = self.env.get_wrapper_attr("sim").is_converged()
         return observation, info
 
 
     def reset(
         self, *, seed: int | None = None, options: dict[str, Any] | None = None
     ) -> tuple[dict[str, Any], dict[str, Any]]:
+        self.prev_action = None
         self.robot.reset()
         if self.home_on_reset:
             self.robot.move_home()
@@ -405,6 +415,7 @@ def __init__(
         else:
             self.robot2world = robot2world
         self.lead_env: gym.Env | None = None
+        self.sim: simulation.Sim | None = None
 
         # make sure all envs are the same type (sim/real)
         for env in self.envs:
@@ -416,6 +427,9 @@ def __init__(
         self._runs_in_sim = self.PLATFORM == RobotPlatform.SIMULATION
         if self._runs_in_sim:
             self._inject_main_greenlet()
+            assert isinstance(self.lead_env, SimEnv), "something is wrong with the env, the base should be type SimEnv"
+            self.sim = self.lead_env.get_wrapper_attr("sim")
+
 
     def _inject_main_greenlet(self):
         main_gr = getcurrent()
@@ -471,9 +485,7 @@ def make_step_gr(env_to_step):
             if self._runs_in_sim:
                 # SIM path: 3. UP: Gather observations
                 # Resume robot greenlet. It returns the step results.
-                res = step_greenlets[key].switch()
-                ob, r, t, tr, i = res
-            
+                ob, r, t, tr, info[key] = step_greenlets[key].switch()
             else:
                 # HARDWARE path
                 act = self._translate_pose(key, action[key], to_world=False)
@@ -510,9 +522,9 @@ def make_reset_gr(env_to_reset, s, o):
                 reset_greenlets[key] = gr
                 gr.switch()
 
-            # SIM path: 2. SIM: reset
+            # SIM path: 2. SIM: apply state from rested wrappers
             assert isinstance(self.lead_env, SimEnv)
-            self.lead_env.reset_sim()
+            self.lead_env.apply_sim_state()
 
 
         for key, env in self.envs.items():

python/rcs/envs/creators.py

@@ -20,6 +20,7 @@
     RelativeTo,
     RobotWrapper,
     SimEnv,
+    CoverWrapper,
 )
 from rcs.envs.sim import (
     GripperWrapperSim,
@@ -128,6 +129,7 @@ def __call__(  # type: ignore
         #     )
         if max_relative_movement is not None:
             env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
+        env = CoverWrapper(env)
 
         return env
 
@@ -186,6 +188,7 @@ def __call__(  # type: ignore
                 BaseCameraSet, SimCameraSet(simulation, cameras, physical_units=True, render_on_demand=True)
             )
             env = CameraSetWrapper(env, camera_set, include_depth=True)
+        env = CoverWrapper(env)
         return env
 
 

python/rcs/rpc/server.py

@@ -26,7 +26,7 @@ def reset(self, **kwargs):
     @rpyc.exposed
     def get_robot_obs(self):
         """Get the current observation using the Wrapper base class if available."""
-        return self.get_wrapper_attr("get_robot_obs")()
+        return self.env.get_wrapper_attr("get_robot_obs")()
 
     @rpyc.exposed
     def unwrapped(self):