refactor(sim envs): simple and lab pick up share more code

- update models
- refactored random cube pos lab wrapper into random cube pos wrapper
- refactored fr3 lab pick up sim success wrapper into pick cube success wrapper
- refactored sim env creation for pick up tasks into own method
- fix cam robot: does not register itself to convergence callback anymore

Author: juelg

.gitignore

@@ -8,7 +8,7 @@ __pycache__
 .ruff_cache
 .mypy_cache
 dist
-.venv
+*venv*
 config.yaml
 MUJOCO_LOG.TXT
 src/pybind/mujoco

CMakeLists.txt

@@ -98,7 +98,7 @@ FetchContent_MakeAvailable(libfranka rl pybind11)
 if (${INCLUDE_UTN_MODELS})
     if (DEFINED GITLAB_MODELS_TOKEN)
         include(download_scenes)
-        set(ref 48c88f05aa34730acbe80bf4b50e7d5a789b927a)
+        set(ref 00869d446940260ec233469182a1ed2290cc60e3)
         FetchContent_Declare(
             scenes
             URL         "https://gitos.rrze.fau.de/api/v4/projects/1100/repository/archive?path=scenes&sha=${ref}"

python/rcsss/envs/factories.py

@@ -21,11 +21,9 @@
 from rcsss.envs.hw import FR3HW
 from rcsss.envs.sim import (
     CollisionGuard,
-    FR3LabPickUpSimSuccessWrapper,
     FR3Sim,
-    FR3SimplePickUpSimSuccessWrapper,
+    PickCubeSuccessWrapper,
     RandomCubePos,
-    RandomCubePosLab,
     SimWrapper,
 )
 from rcsss.envs.space_utils import Vec7Type
@@ -211,11 +209,37 @@ def __call__(  # type: ignore
         )
 
 
+def make_sim_task_envs(
+    mjcf: str,
+    render_mode: str = "human",
+    control_mode: ControlMode = ControlMode.CARTESIAN_TRPY,
+    delta_actions: bool = True,
+    camera_cfg: SimCameraConfig | None = None,
+) -> gym.Env:
+
+    env_rel = fr3_sim_env(
+        control_mode=control_mode,
+        robot_cfg=default_fr3_sim_robot_cfg(mjcf),
+        collision_guard=False,
+        gripper_cfg=default_fr3_sim_gripper_cfg(),
+        camera_set_cfg=camera_cfg,
+        max_relative_movement=(0.2, np.deg2rad(45)) if delta_actions else None,
+        relative_to=RelativeTo.LAST_STEP,
+        mjcf=mjcf,
+        sim_wrapper=RandomCubePos,
+    )
+    env_rel = PickCubeSuccessWrapper(env_rel)
+    if render_mode == "human":
+        env_rel.get_wrapper_attr("sim").open_gui()
+
+    return env_rel
+
+
 class FR3SimplePickUpSim(EnvCreator):
     def __call__(  # type: ignore
         self,
         render_mode: str = "human",
-        control_mode: str = "xyzrpy",
+        control_mode: ControlMode = ControlMode.CARTESIAN_TRPY,
         resolution: tuple[int, int] | None = None,
         frame_rate: int = 10,
         delta_actions: bool = True,
@@ -239,33 +263,14 @@ def __call__(  # type: ignore
             frame_rate=frame_rate,
             physical_units=True,
         )
-        env_rel = fr3_sim_env(
-            control_mode=(
-                ControlMode.CARTESIAN_TRPY
-                if control_mode == "xyzrpy"
-                else ControlMode.JOINTS if control_mode == "joints" else ControlMode.CARTESIAN_TQuart
-            ),
-            robot_cfg=default_fr3_sim_robot_cfg(),
-            collision_guard=False,
-            gripper_cfg=default_fr3_sim_gripper_cfg(),
-            camera_set_cfg=camera_cfg,
-            max_relative_movement=(0.2, np.deg2rad(45)) if delta_actions else None,
-            relative_to=RelativeTo.LAST_STEP,
-            mjcf="fr3_simple_pick_up",
-            sim_wrapper=RandomCubePos,
-        )
-        env_rel = FR3SimplePickUpSimSuccessWrapper(env_rel)
-        if render_mode == "human":
-            env_rel.get_wrapper_attr("sim").open_gui()
-
-        return env_rel
+        return make_sim_task_envs("simple_pick_up", render_mode, control_mode, delta_actions, camera_cfg)
 
 
 class FR3SimplePickUpSimDigitHand(EnvCreator):
     def __call__(  # type: ignore
         self,
         render_mode: str = "human",
-        control_mode: str = "xyzrpy",
+        control_mode: ControlMode = ControlMode.CARTESIAN_TRPY,
         resolution: tuple[int, int] | None = None,
         frame_rate: int = 10,
         delta_actions: bool = True,
@@ -282,27 +287,7 @@ def __call__(  # type: ignore
             frame_rate=frame_rate,
             physical_units=True,
         )
-
-        env_rel = fr3_sim_env(
-            control_mode=(
-                ControlMode.CARTESIAN_TRPY
-                if control_mode == "xyzrpy"
-                else ControlMode.JOINTS if control_mode == "joints" else ControlMode.CARTESIAN_TQuart
-            ),
-            robot_cfg=default_fr3_sim_robot_cfg("fr3_simple_pick_up_digit_hand"),
-            collision_guard=False,
-            gripper_cfg=default_fr3_sim_gripper_cfg(),
-            camera_set_cfg=camera_cfg,
-            max_relative_movement=(0.2, np.deg2rad(45)) if delta_actions else None,
-            relative_to=RelativeTo.LAST_STEP,
-            mjcf="fr3_simple_pick_up_digit_hand",
-            sim_wrapper=RandomCubePos,
-        )
-        env_rel = FR3SimplePickUpSimSuccessWrapper(env_rel)
-        if render_mode == "human":
-            env_rel.get_wrapper_attr("sim").open_gui()
-
-        return env_rel
+        return make_sim_task_envs("fr3_simple_pick_up_digit_hand", render_mode, control_mode, delta_actions, camera_cfg)
 
 
 class CamRobot(gym.Wrapper):
@@ -312,7 +297,7 @@ def __init__(self, env, cam_robot_joints: Vec7Type):
         self.unwrapped: FR3Env
         assert isinstance(self.unwrapped.robot, sim.FR3), "Robot must be a sim.FR3 instance."
         self.sim = env.get_wrapper_attr("sim")
-        self.cam_robot = rcsss.sim.FR3(self.sim, "1", env.unwrapped.robot.get_ik())
+        self.cam_robot = rcsss.sim.FR3(self.sim, "1", env.unwrapped.robot.get_ik(), register_convergence_callback=False)
         self.cam_robot.set_parameters(default_fr3_sim_robot_cfg("lab_simple_pick_up_digit_hand"))
         self.cam_robot_joints = cam_robot_joints
 
@@ -332,7 +317,7 @@ def __call__(  # type: ignore
         self,
         cam_robot_joints: Vec7Type,
         render_mode: str = "human",
-        control_mode: str = "xyzrpy",
+        control_mode: ControlMode = ControlMode.CARTESIAN_TRPY,
         resolution: tuple[int, int] | None = None,
         frame_rate: int = 10,
         delta_actions: bool = True,
@@ -352,25 +337,11 @@ def __call__(  # type: ignore
             frame_rate=frame_rate,
             physical_units=True,
         )
-
-        env_rel = fr3_sim_env(
-            control_mode=(
-                ControlMode.CARTESIAN_TRPY
-                if control_mode == "xyzrpy"
-                else ControlMode.JOINTS if control_mode == "joints" else ControlMode.CARTESIAN_TQuart
-            ),
-            robot_cfg=default_fr3_sim_robot_cfg("lab_simple_pick_up_digit_hand"),
-            collision_guard=False,
-            gripper_cfg=default_fr3_sim_gripper_cfg(),
-            camera_set_cfg=camera_cfg,
-            max_relative_movement=(0.2, np.deg2rad(45)) if delta_actions else None,
-            relative_to=RelativeTo.LAST_STEP,
-            mjcf="lab_simple_pick_up_digit_hand",
-            sim_wrapper=RandomCubePosLab,
+        env_rel = make_sim_task_envs(
+            "lab_simple_pick_up_digit_hand",
+            render_mode,
+            control_mode,
+            delta_actions,
+            camera_cfg,
         )
-        env_rel = FR3LabPickUpSimSuccessWrapper(env_rel)
-        env_rel = CamRobot(env_rel, cam_robot_joints)
-        if render_mode == "human":
-            sim = env_rel.get_wrapper_attr("sim")
-            sim.open_gui()
-        return env_rel
+        return CamRobot(env_rel, cam_robot_joints)

python/rcsss/envs/sim.py

@@ -176,7 +176,7 @@ def env_from_xml_paths(
         )
 
 
-class RandomCubePosLab(SimWrapper):
+class RandomCubePos(SimWrapper):
     """Wrapper to randomly place cube in the lab environments."""
 
     def reset(
@@ -197,25 +197,7 @@ def reset(
         return obs, info
 
 
-class RandomCubePos(SimWrapper):
-    """Wrapper to randomly place cube in the FR3SimplePickUpSim environment."""
-
-    def reset(
-        self, seed: int | None = None, options: dict[str, Any] | None = None
-    ) -> tuple[dict[str, Any], dict[str, Any]]:
-        obs, info = super().reset(seed=seed, options=options)
-        self.sim.step(1)
-        iso_cube = [0.498, 0.0, 0.226]
-        pos_z = 0.03
-        pos_x = iso_cube[0] + np.random.random() * 0.2 - 0.1
-        pos_y = iso_cube[1] + np.random.random() * 0.2 - 0.1
-
-        self.sim.data.joint("yellow-box-joint").qpos[:3] = [pos_x, pos_y, pos_z]
-
-        return obs, info
-
-
-class FR3SimplePickUpSimSuccessWrapper(gym.Wrapper):
+class PickCubeSuccessWrapper(gym.Wrapper):
     """Wrapper to check if the cube is successfully picked up in the FR3SimplePickUpSim environment."""
 
     EE_HOME = np.array([0.34169773, 0.00047028, 0.4309004])
@@ -241,34 +223,3 @@ def step(self, action: dict[str, Any]):
         reward = -diff_cube_home - diff_ee_cube
 
         return obs, reward, success, truncated, info
-
-
-class FR3LabPickUpSimSuccessWrapper(gym.Wrapper):
-    """
-    Wrapper to check if the cube is successfully picked up in the LabPickupSim environments.
-    This is also used to handle the reset of the second robot.
-    """
-
-    EE_HOME = np.array([0.34169773, 0.00047028, 0.4309004])
-
-    def __init__(self, env):
-        super().__init__(env)
-        self.unwrapped: FR3Env
-        assert isinstance(self.unwrapped.robot, sim.FR3), "Robot must be a sim.FR3 instance."
-        self.sim = env.get_wrapper_attr("sim")
-
-    def step(self, action: dict[str, Any]):
-        obs, reward, done, truncated, info = super().step(action)
-
-        success = (
-            self.sim.data.joint("yellow-box-joint").qpos[2] > 0.3
-            and obs["gripper"] == GripperWrapper.BINARY_GRIPPER_CLOSED
-        )
-        diff_ee_cube = np.linalg.norm(
-            self.sim.data.joint("yellow-box-joint").qpos[:3]
-            - self.unwrapped.robot.get_cartesian_position().translation()
-        )
-        diff_cube_home = np.linalg.norm(self.sim.data.joint("yellow-box-joint").qpos[:3] - self.EE_HOME)
-        reward = -diff_cube_home - diff_ee_cube
-
-        return obs, reward, success, truncated, info