Jin/taxim cleanup

examples/fr3/grasp_digit_demo.py → examples/taxim/grasp_digit_demo.py

@@ -7,70 +7,75 @@
 from rcs._core.common import Pose
 from rcs._core.sim import SimRobot
 from rcs.envs.base import GripperWrapper
-from rcs_tacto.creators import FR3TactoSimplePickUpSimEnvCreator
-from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 
+def _progress(iterable):
+    try:
+        from tqdm import tqdm
+    except ImportError:
+        return iterable
+    return tqdm(iterable)
+
+
 class PickUpDemo:
     def __init__(self, env: gym.Env):
         self.env = env
-        self._robot = cast(SimRobot, self.env.get_wrapper_attr("robot"))
+        self._robot = cast(SimRobot, self.env.get_wrapper_attr("robot")["robot"])
         self.home_pose = self._robot.get_cartesian_position()
 
     def _action(self, pose: Pose, gripper: list[float]) -> dict[str, Any]:
-        return {"xyzrpy": pose.xyzrpy(), "gripper": gripper}
+        return {"robot": {"xyzrpy": pose.xyzrpy(), "gripper": gripper}}
 
-    def get_object_pose(self, geom_name) -> Pose:
+    def get_object_pose(self, geom_name: str) -> 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)
-        )
+        ) * Pose(rpy_vector=np.array([0, 0, -np.pi / 4]), translation=np.array([0.0, 0.0, 0.0]))
         return self._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)
-            waypoints.append(start_pose.interpolate(end_pose, t))
-        return waypoints
+        return [start_pose.interpolate(end_pose, i / num_waypoints) for i in range(num_waypoints + 1)]
 
-    def step(self, action: dict) -> dict:
+    def step(self, action: dict[str, Any]) -> dict[str, Any]:
         return self.env.step(action)[0]
 
-    def plan_linear_motion(self, geom_name: str, delta_up: float, num_waypoints: int = 200) -> list[Pose]:
+    def plan_linear_motion(self, geom_name: str, delta_up: float, num_waypoints: int = 20) -> list[Pose]:
         end_eff_pose = self._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]))  # type: ignore
+        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: list[float] = GripperWrapper.BINARY_GRIPPER_OPEN) -> dict:
-        for i in range(len(waypoints)):
-            obs = self.step(self._action(waypoints[i], gripper))
+    def execute_motion(
+        self, waypoints: list[Pose], gripper: list[float] = GripperWrapper.BINARY_GRIPPER_OPEN
+    ) -> dict[str, Any]:
+        obs: dict[str, Any] = {}
+        for waypoint in waypoints:
+            obs = self.step(self._action(waypoint, gripper))
         return obs
 
     def approach(self, geom_name: str):
-        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.2, num_waypoints=5)
+        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.2, num_waypoints=60)
         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.01, num_waypoints=15)
+        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.09, num_waypoints=60)
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_OPEN)
 
-        self.step(self._action(Pose(), GripperWrapper.BINARY_GRIPPER_CLOSED))
+        for _ in range(4):
+            self.step(self._action(self._robot.get_cartesian_position(), GripperWrapper.BINARY_GRIPPER_CLOSED))
 
-        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.2, num_waypoints=15)
+        waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.2, num_waypoints=60)
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
 
     def move_home(self):
         end_eff_pose = self._robot.get_cartesian_position()
-        waypoints = self.generate_waypoints(end_eff_pose, self.home_pose, num_waypoints=10)
+        waypoints = self.generate_waypoints(end_eff_pose, self.home_pose, num_waypoints=60)
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
 
     def pickup(self, geom_name: str):
@@ -80,17 +85,19 @@ def pickup(self, geom_name: str):
 
 
 def main():
-    env_fact = FR3TactoSimplePickUpSimEnvCreator()
-    env = env_fact(
-        render_mode="human",
-        delta_actions=False,
-    )
-
-    for _ in tqdm(range(100)):
-        # reset the environment
+    try:
+        from rcs_taxim.creators import FR3TaximSimplePickUpSimEnvCreator
+    except ImportError as exc:
+        msg = "This example requires the rcs_taxim extension, install it with `pip install -e extensions/rcs_taxim`."
+        raise ImportError(msg) from exc
+
+    env_fact = FR3TaximSimplePickUpSimEnvCreator()
+    env = env_fact(render_mode="human", delta_actions=False)
+
+    for _ in _progress(range(100)):
         env.reset()
         controller = PickUpDemo(env)
-        controller.pickup("yellow_box_geom")
+        controller.pickup("_box_geom")
     env.close()
 
 

extensions/rcs_taxim/README.md

@@ -0,0 +1,11 @@
+# RCS Taxim Extension
+
+This extension provides integration with the [Taxim](https://github.com/Robo-Touch/Taxim) tactile sensor simulator.
+
+It ports the tactile pick-up example onto the current RCS scene/config stack by swapping in a TAXIM-equipped Franka hand while reusing the current FR3 pick setup.
+
+## Installation
+
+```shell
+pip install -ve extensions/rcs_taxim
+```

extensions/rcs_taxim/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "rcs_taxim"
+version = "0.7.2"
+description = "RCS integration of mujoco-taxim"
+dependencies = [
+  "rcs>=0.7.2",
+  "omegaconf",
+  "mujoco-taxim@git+https://github.com/utn-air/mujoco-taxim.git@norm2tex",
+]
+readme = "README.md"
+maintainers = [
+  { name = "Tobias Jülg", email = "tobias.juelg@utn.de" },
+  { name = "Seongjin Bien", email = "seongjin.bien@utn.de" },
+]
+authors = [{ name = "Seongjin Bien", email = "seongjin.bien@utn.de" }]
+requires-python = ">=3.10"
+
+[tool.black]
+line-length = 120
+target-version = ["py310"]
+
+[tool.isort]
+profile = "black"

extensions/rcs_taxim/src/rcs_taxim/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.7.2"

extensions/rcs_taxim/src/rcs_taxim/creators.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+from rcs._core.common import GripperType, Pose
+from rcs._core.sim import SimCameraConfig, SimGripperConfig
+from rcs.envs.base import ControlMode, RelativeTo
+from rcs.envs.configs import EmptyWorldFR3
+from rcs_taxim.taxim_wrapper import TaximSimWrapper, _robotiq2f85_digit_model_path
+
+import rcs
+
+_TAXIM_GRIPPER_TYPE = GripperType("Robotiq2F85Digit")
+
+
+def _prefixed(name: str) -> str:
+    return f"gripper{name}"
+
+
+def _make_camera_cfgs(
+    base_cfgs: dict[str, SimCameraConfig],
+    cam_list: tuple[str, ...],
+    resolution: tuple[int, int] | None,
+    frame_rate: int,
+) -> dict[str, SimCameraConfig]:
+    camera_cfgs: dict[str, SimCameraConfig] = {}
+    for cam in cam_list:
+        if cam not in base_cfgs:
+            available = ", ".join(sorted(base_cfgs))
+            msg = f"Unknown camera {cam!r}. Available cameras: {available}"
+            raise ValueError(msg)
+        cfg = copy.deepcopy(base_cfgs[cam])
+        if resolution is not None:
+            cfg.resolution_width = resolution[0]
+            cfg.resolution_height = resolution[1]
+        if frame_rate > 0:
+            cfg.frame_rate = frame_rate
+        camera_cfgs[cam] = cfg
+    return camera_cfgs
+
+
+def _taxim_gripper_cfg() -> SimGripperConfig:
+    return SimGripperConfig(
+        epsilon_inner=0.005,
+        epsilon_outer=0.005,
+        seconds_between_callbacks=0.1,
+        ignored_collision_geoms=[],
+        collision_geoms=[],
+        collision_geoms_fingers=[],
+        joints=["right_driver_joint", "left_driver_joint"],
+        max_joint_width=0.005,
+        min_joint_width=1.0,
+        actuator="fingers_actuator",
+        max_actuator_width=0,
+        min_actuator_width=255,
+        gripper_type=_TAXIM_GRIPPER_TYPE,
+    )
+
+
+class FR3TaximSimplePickUpSimEnvCreator:
+    def __call__(
+        self,
+        render_mode: str = "human",
+        control_mode: ControlMode = ControlMode.CARTESIAN_TRPY,
+        resolution: tuple[int, int] | None = None,
+        frame_rate: int = 0,
+        delta_actions: bool = True,
+        cam_list: tuple[str, ...] | None = None,
+        taxim_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> gym.Env:
+        binary_gripper = kwargs.pop("binary_gripper", True)
+        home_on_reset = kwargs.pop("home_on_reset", True)
+        max_relative_movement = kwargs.pop("max_relative_movement", None)
+        if kwargs:
+            unexpected = ", ".join(sorted(kwargs))
+            msg = f"Unexpected keyword arguments: {unexpected}"
+            raise TypeError(msg)
+
+        rcs.GRIPPER_PATHS[_TAXIM_GRIPPER_TYPE] = str(_robotiq2f85_digit_model_path())
+
+        scene = EmptyWorldFR3()
+        cfg = scene.config()
+        cfg.control_mode = control_mode
+        cfg.headless = render_mode != "human"
+        cfg.sim_cfg.realtime = render_mode == "human"
+        cfg.sim_cfg.async_control = render_mode == "human"
+        cfg.sim_cfg.frequency = 30
+        cfg.wrapper_cfg.binary_gripper = binary_gripper
+        cfg.wrapper_cfg.home_on_reset = home_on_reset
+        cfg.max_relative_movement = (
+            max_relative_movement if max_relative_movement is not None else (0.2, np.deg2rad(45))
+        )
+        cfg.relative_to = RelativeTo.LAST_STEP if delta_actions else RelativeTo.NONE
+        if not delta_actions:
+            cfg.max_relative_movement = None
+        cfg.gripper_cfgs = {"robot": _taxim_gripper_cfg()}
+        cfg.gripper_offsets = None
+        cfg.root_frame_objects = {
+            "": (
+                rcs.OBJECT_PATHS["green_cube"],
+                Pose(translation=np.array([0.5, 0.0, 0.05]), quaternion=np.array([0.0, 0.0, 0.0, 1.0])),
+            )
+        }
+
+        cam_list = tuple(cam_list or ())
+        if cam_list:
+            cfg.camera_cfgs = _make_camera_cfgs(cfg.camera_cfgs or {}, cam_list, resolution, frame_rate)
+            cfg.camera_adds = (
+                {name: cfg.camera_adds[name] for name in cam_list} if cfg.camera_adds is not None else None
+            )
+        else:
+            cfg.camera_cfgs = None
+            cfg.camera_adds = None
+
+        env = scene.create_env(cfg)
+        merged_taxim_kwargs: dict[str, Any] = {
+            "taxim_sites": [_prefixed("left_digit_pad"), _prefixed("right_digit_pad")],
+            "taxim_pad_geoms": [_prefixed("left_digit_pad"), _prefixed("right_digit_pad")],
+            "target_geom_mesh_dict": {"_box_geom": "_box_geom"},
+            "taxim_sensor_type": "digit",
+            "taxim_fps": 60,
+            "enable_depth": True,
+            "visualize": True,
+        }
+        if taxim_kwargs is not None:
+            merged_taxim_kwargs.update(taxim_kwargs)
+        env = TaximSimWrapper(env, **merged_taxim_kwargs)
+        if render_mode == "human":
+            env.get_wrapper_attr("sim").open_gui()
+        return env