feat(sim): support for camera robot lab scene

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 b8234983a5e35e222c955f9145ac4cd129827a8e)
+        set(ref 377dffded50e7815fa65e84ab4b0eb3198065881)
         FetchContent_Declare(
             scenes
             URL         "https://gitos.rrze.fau.de/api/v4/projects/1100/repository/archive?path=scenes&sha=${ref}"

python/examples/grasp_demo.py

@@ -32,6 +32,7 @@ def generate_waypoints(self, start_pose: Pose, end_pose: Pose, num_waypoints: in
         for i in range(num_waypoints + 1):
             t = i / (num_waypoints + 1)
             waypoints.append(start_pose.interpolate(end_pose, t))
+        waypoints.append(end_pose)
         return waypoints
 
     def step(self, action: dict) -> dict:
@@ -83,6 +84,8 @@ def pickup(self, geom_name: str):
 def main():
     env = gym.make("rcs/SimplePickUpSim-v0", render_mode="human", delta_actions=True)
     env.reset()
+    print(env.unwrapped.robot.get_cartesian_position().translation())  # type: ignore
+    # assert False
     controller = PickUpDemo(env)
     controller.pickup("yellow_box_geom")
 

python/examples/grasp_demo_lab.py

@@ -0,0 +1,116 @@
+import logging
+import sys
+from typing import Any, cast
+
+import gymnasium as gym
+import mujoco
+import numpy as np
+from rcsss._core.common import Pose
+import rcsss.envs.base
+from rcsss.envs.base import FR3Env, GripperWrapper
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+class PickUpDemo:
+    def __init__(self, env: gym.Env):
+        self.env = env
+        self.unwrapped: FR3Env = cast(FR3Env, self.env.unwrapped)
+        self.home_pose = self.unwrapped.robot.get_cartesian_position()
+
+    def _action(self, pose: Pose, gripper: float) -> dict[str, Any]:
+        return {"xyzrpy": pose.xyzrpy(), "gripper": gripper}
+
+    def get_object_pose(self, geom_name) -> Pose:
+        model = self.env.get_wrapper_attr("sim").model
+        data = self.env.get_wrapper_attr("sim").data
+
+        geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, geom_name)
+        obj_pose_world_coordinates = Pose(
+            translation=data.geom_xpos[geom_id], rotation=data.geom_xmat[geom_id].reshape(3, 3)
+        )
+        return self.env.unwrapped.robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates)
+
+    def generate_waypoints(self, start_pose: Pose, end_pose: Pose, num_waypoints: int) -> list[Pose]:
+        waypoints = []
+        for i in range(num_waypoints + 1):
+            t = i / (num_waypoints + 1)
+            waypoints.append(start_pose.interpolate(end_pose, t))
+        return waypoints
+
+    def step(self, action: np.ndarray) -> dict:
+        return self.env.step(action)
+
+    def plan_linear_motion(self, geom_name: str, delta_up: float, num_waypoints: int = 200) -> list[Pose]:
+        end_eff_pose = self.unwrapped.robot.get_cartesian_position()
+        goal_pose = self.get_object_pose(geom_name=geom_name)
+        goal_pose *= Pose(translation=np.array([0, 0, delta_up]), quaternion=np.array([1, 0, 0, 0]))
+        return self.generate_waypoints(end_eff_pose, goal_pose, num_waypoints=num_waypoints)
+
+    def execute_motion(self, waypoints: list[Pose], gripper: float = GripperWrapper.BINARY_GRIPPER_OPEN) -> dict:
+        for i in range(1, len(waypoints)):
+            # calculate delta action
+            # delta_action = waypoints[i] * waypoints[i - 1].inverse()
+            pose = rcsss.common.Pose(translation=waypoints[i].translation(),
+                                     rpy_vector=np.array([-3.14159265e+00,  1.57009246e-16, 0]))
+
+            # act = self._action(delta_action, gripper)
+            act = self._action(pose, gripper)
+
+            obs = self.step(act)
+            ik_success = obs[-1]["ik_success"]
+            if not obs[-1]["ik_success"]:
+                trans_source, rot_source = waypoints[i - 1].translation(), waypoints[i - 1].rotation_rpy().as_vector()
+                trans_dest, rot_des = waypoints[i].translation(), waypoints[i].rotation_rpy().as_vector()
+                msg = (f"ik success: {ik_success} when attempting to move from trans: {trans_source}, rot: {rot_source}\n "
+                       f"to trans: {trans_dest} rot: {rot_des}!")
+                logger.warning(msg)
+                assert False, msg
+        return obs
+
+    def approach(self, geom_name: str):
+        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.2, num_waypoints=50)
+        self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_OPEN)
+
+    def grasp(self, geom_name: str):
+
+        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0, num_waypoints=50)
+        self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_OPEN)
+
+        self.step(self._action(Pose(), GripperWrapper.BINARY_GRIPPER_CLOSED))
+
+        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.2, num_waypoints=50)
+        self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
+
+    def move_home(self):
+        end_eff_pose = self.unwrapped.robot.get_cartesian_position()
+        waypoints = self.generate_waypoints(end_eff_pose, self.home_pose, num_waypoints=10)
+        self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
+
+    def pickup(self, geom_name: str):
+        self.approach(geom_name)
+        self.grasp(geom_name)
+        self.move_home()
+
+
+def main():
+    # env = gym.make(
+    #     "rcs/SimplePickUpSimDigitHand-v0",
+    #     render_mode="human",
+    #     delta_actions=True
+    # )
+
+    env = gym.make(
+        "rcs/FR3LabPickUpSimDigitHand-v0",
+        render_mode="human",
+        delta_actions=False,
+        robot2_cam_pose=[0.1243549, -1.4711298, 1.2246249, -1.9944441, 0.0872650, 1.3396115, 2.1275465],
+    )
+    env.reset()
+    controller = PickUpDemo(env)
+    controller.pickup("yellow_box_geom")
+
+
+if __name__ == "__main__":
+    main()

python/rcsss/__init__.py

@@ -6,7 +6,12 @@
 from gymnasium import register
 from rcsss import camera, control, envs, sim
 from rcsss._core import __version__, common, hw
-from rcsss.envs.factories import FR3Real, FR3SimplePickUpSim
+from rcsss.envs.factories import (
+    FR3LabPickUpSimDigitHand,
+    FR3Real,
+    FR3SimplePickUpSim,
+    FR3SimplePickUpSimDigitHand,
+)
 
 # available mujoco scenes
 scenes = {
@@ -21,6 +26,16 @@
     id="rcs/SimplePickUpSim-v0",
     entry_point=FR3SimplePickUpSim(),
 )
+
+register(
+    id="rcs/SimplePickUpSimDigitHand-v0",
+    entry_point=FR3SimplePickUpSimDigitHand(),
+)
+register(
+    id="rcs/FR3LabPickUpSimDigitHand-v0",
+    entry_point=FR3LabPickUpSimDigitHand(),
+)
+
 register(
     id="rcs/FR3Real-v0",
     entry_point=FR3Real(),

python/rcsss/envs/factories.py

@@ -27,8 +27,10 @@
     FR3Sim,
     FR3SimplePickUpSimSuccessWrapper,
     RandomCubePos,
-    SimWrapper,
+    RandomCubePosLab,
+    SimWrapper, FR3LabPickUpSimSuccessWrapper,
 )
+from rcsss.envs.utils import get_urdf_path, set_tcp_offset, default_fr3_sim_robot_cfg
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -66,19 +68,6 @@ def default_realsense(name2id: dict[str, str] | None) -> RealSenseCameraSet | No
     return RealSenseCameraSet(cam_cfg)
 
 
-def get_urdf_path(urdf_path: str | PathLike | None, allow_none_if_not_found: bool = False) -> str | None:
-    if urdf_path is None and "lab" in rcsss.scenes:
-        urdf_path = rcsss.scenes["lab"].parent / "fr3.urdf"
-        assert urdf_path.exists(), "Automatic deduced urdf path does not exist. Corrupted models directory."
-        logger.info("Using automatic found urdf.")
-    elif urdf_path is None and not allow_none_if_not_found:
-        msg = "This pip package was not built with the UTN lab models, please pass the urdf and mjcf path to use simulation or collision guard."
-        raise ValueError(msg)
-    elif urdf_path is None:
-        logger.warning("No urdf path was found. Proceeding, but set_cartesian methods will result in errors.")
-    return str(urdf_path) if urdf_path is not None else None
-
-
 def fr3_hw_env(
     ip: str,
     control_mode: ControlMode,
@@ -149,13 +138,6 @@ def fr3_hw_env(
     return env
 
 
-def default_fr3_sim_robot_cfg():
-    cfg = sim.FR3Config()
-    cfg.tcp_offset = rcsss.common.Pose(rcsss.common.FrankaHandTCPOffset())
-    cfg.realtime = False
-    return cfg
-
-
 def default_fr3_sim_gripper_cfg():
     return sim.FHConfig()
 
@@ -211,8 +193,10 @@ def fr3_sim_env(
     assert urdf_path is not None
     if mjcf not in rcsss.scenes:
         logger.warning("mjcf not found as key in scenes, interpreting mjcf as path the mujoco scene xml")
-
-    simulation = sim.Sim(rcsss.scenes.get(str(mjcf), mjcf))
+    mjb_file = rcsss.scenes.get(str(mjcf), mjcf)
+    simulation = sim.Sim(mjb_file)
+    # setting the tcp offset
+    set_tcp_offset(robot_cfg, simulation)
     ik = rcsss.common.IK(urdf_path)
     robot = rcsss.sim.FR3(simulation, "0", ik)
     robot.set_parameters(robot_cfg)
@@ -319,3 +303,96 @@ def __call__(  # type: ignore
             env_rel.get_wrapper_attr("sim").open_gui()
 
         return env_rel
+
+
+class FR3SimplePickUpSimDigitHand(EnvCreator):
+    def __call__(  # type: ignore
+        self,
+        render_mode: str = "human",
+        control_mode: str = "xyzrpy",
+        resolution: tuple[int, int] | None = None,
+        frame_rate: int = 10,
+        delta_actions: bool = True,
+    ) -> gym.Env:
+        if resolution is None:
+            resolution = (256, 256)
+
+        cameras = {"eye-in-hand_0": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed))}
+
+        camera_cfg = SimCameraSetConfig(
+            cameras=cameras,
+            resolution_width=resolution[0],
+            resolution_height=resolution[1],
+            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_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
+
+
+class FR3LabPickUpSimDigitHand(EnvCreator):
+    def __call__(  # type: ignore
+        self,
+        render_mode: str = "human",
+        control_mode: str = "xyzrpy",
+        resolution: tuple[int, int] | None = None,
+        frame_rate: int = 10,
+        delta_actions: bool = True,
+        robot2_cam_pose: list[int] | None = None,
+    ) -> gym.Env:
+        if resolution is None:
+            resolution = (256, 256)
+
+        cameras = {
+            "eye-in-hand_0": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),
+            "eye-in-hand_1": SimCameraConfig(identifier="eye-in-hand_1", type=int(CameraType.fixed)),
+        }
+
+        camera_cfg = SimCameraSetConfig(
+            cameras=cameras,
+            resolution_width=resolution[0],
+            resolution_height=resolution[1],
+            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="lab_simple_pick_up_digit_hand",
+            sim_wrapper=RandomCubePosLab,
+        )
+        env_rel = FR3LabPickUpSimSuccessWrapper(env_rel, robot2_cam_pose=robot2_cam_pose)
+        sim = env_rel.get_wrapper_attr("sim")
+        if render_mode == "human":
+            sim.open_gui()
+
+        return env_rel