diff --git a/examples/fr3/grasp_digit_demo.py b/examples/taxim/grasp_digit_demo.py similarity index 61% rename from examples/fr3/grasp_digit_demo.py rename to examples/taxim/grasp_digit_demo.py index c90fccad..c1f29ca7 100644 --- a/examples/fr3/grasp_digit_demo.py +++ b/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() diff --git a/extensions/rcs_taxim/README.md b/extensions/rcs_taxim/README.md new file mode 100644 index 00000000..2625c78d --- /dev/null +++ b/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 +``` diff --git a/extensions/rcs_taxim/pyproject.toml b/extensions/rcs_taxim/pyproject.toml new file mode 100644 index 00000000..15e7a6e6 --- /dev/null +++ b/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" diff --git a/extensions/rcs_taxim/src/rcs_taxim/__init__.py b/extensions/rcs_taxim/src/rcs_taxim/__init__.py new file mode 100644 index 00000000..bc8c296f --- /dev/null +++ b/extensions/rcs_taxim/src/rcs_taxim/__init__.py @@ -0,0 +1 @@ +__version__ = "0.7.2" diff --git a/extensions/rcs_taxim/src/rcs_taxim/creators.py b/extensions/rcs_taxim/src/rcs_taxim/creators.py new file mode 100644 index 00000000..3ebd54f0 --- /dev/null +++ b/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 diff --git a/extensions/rcs_taxim/src/rcs_taxim/taxim_wrapper.py b/extensions/rcs_taxim/src/rcs_taxim/taxim_wrapper.py new file mode 100644 index 00000000..5915ba62 --- /dev/null +++ b/extensions/rcs_taxim/src/rcs_taxim/taxim_wrapper.py @@ -0,0 +1,132 @@ +import copy +from pathlib import Path +from typing import Any + +import gymnasium as gym +import numpy as np +import TaximSensor + + +def _robotiq2f85_digit_model_path() -> str: + return str(Path(TaximSensor.__file__).parent.parent.parent / "assets/robotiq_2f85/robotiq_2f85.xml") + + +class TaximSimWrapper(gym.Wrapper): + """Wrapper to render TAXIM tactile observations alongside regular RCS observations.""" + + def __init__( + self, + env: gym.Env, + taxim_sites: list[str], + taxim_pad_geoms: list[str], + target_geom_mesh_dict: dict[str, str], + target_geom_normal_map_dict: dict[str, str] | None = None, + taxim_sensor_type: str = "digit", + taxim_bg_idx: int = 0, + taxim_bg_randomize: bool = False, + enable_depth: bool = False, + taxim_fps: int = 60, + visualize: bool = False, + ): + super().__init__(env) + self.taxim_sensors: list[Any] = [] + self.model = self.env.get_wrapper_attr("sim").model + self.data = self.env.get_wrapper_attr("sim").data + self.last_tactile_frames: dict[str, dict[str, dict[str, Any]]] = {} + + self.taxim_sites = taxim_sites + self.taxim_pad_geoms = taxim_pad_geoms + self.target_geom_mesh_dict = target_geom_mesh_dict + self.target_geom_normal_map_dict = target_geom_normal_map_dict + self.taxim_sensor_type = taxim_sensor_type + self.taxim_bg_idx = taxim_bg_idx + self.taxim_bg_randomize = taxim_bg_randomize + self.taxim_fps = taxim_fps + self.taxim_last_render = -1.0 + self.enable_depth = enable_depth + self.initialized = False + self.visualize = visualize + + self.observation_space = copy.deepcopy(env.observation_space) + if not isinstance(self.observation_space, gym.spaces.Dict): + msg = "Expected wrapped observation space to be a gym.spaces.Dict." + raise TypeError(msg) + + frame_spaces = self.observation_space.spaces.get("frames") + if frame_spaces is None: + frame_spaces = gym.spaces.Dict({}) + self.observation_space.spaces["frames"] = frame_spaces + if not isinstance(frame_spaces, gym.spaces.Dict): + msg = "Expected frames observation space to be a gym.spaces.Dict." + raise TypeError(msg) + + tactile_space: dict[str, gym.Space[Any]] = { + "rgb": gym.spaces.Dict( + { + "data": gym.spaces.Box(low=0, high=255, shape=(320, 240, 3), dtype=np.uint8), + } + ) + } + if self.enable_depth: + tactile_space["depth"] = gym.spaces.Dict( + { + "data": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(320, 240), dtype=np.float64), + } + ) + frame_spaces.spaces.update( + {f"tactile_{site}": gym.spaces.Dict(copy.deepcopy(tactile_space)) for site in self.taxim_sites} + ) + + def _ensure_initialized(self) -> None: + if self.initialized: + return + + for site, pad_geom in zip(self.taxim_sites, self.taxim_pad_geoms, strict=True): + sensor = TaximSensor.TaximSensor(resize=(240, 320), sensor_type=self.taxim_sensor_type, preprocess_bg=False) + sensor.add_camera_mujoco(site, self.model, self.data) + sensor.change_bg(self.taxim_bg_idx) + for geom, mesh in self.target_geom_mesh_dict.items(): + normal_map_path = None + if self.target_geom_normal_map_dict is not None and geom in self.target_geom_normal_map_dict: + normal_map_path = self.target_geom_normal_map_dict[geom] + sensor.add_geom_mujoco(geom, self.model, self.data, mesh, normal_map_path=normal_map_path) + sensor.set_sensor_pad_geom(pad_geom) + self.taxim_sensors.append(sensor) + self.initialized = True + + def _render_tactile_frames(self, visualize: bool) -> dict[str, dict[str, dict[str, Any]]]: + frames: dict[str, dict[str, dict[str, Any]]] = {} + for site, sensor in zip(self.taxim_sites, self.taxim_sensors, strict=True): + rgb, depth, _ = sensor.render_taxim( + self.model, self.data, visualize=self.visualize, cycle_bg=self.taxim_bg_randomize + ) + tactile_obs: dict[str, dict[str, Any]] = {"rgb": {"data": rgb}} + if self.enable_depth: + tactile_obs["depth"] = {"data": depth} + frames[f"tactile_{site}"] = tactile_obs + return frames + + def _update_obs(self, obs: dict[str, Any]) -> None: + frames = obs.setdefault("frames", {}) + for site, tactile_obs in self.last_tactile_frames.items(): + frames[site] = copy.deepcopy(tactile_obs) + + 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._ensure_initialized() + self.taxim_last_render = -1.0 + self.last_tactile_frames = self._render_tactile_frames(visualize=False) + self._update_obs(obs) + return obs, info + + def step(self, action: dict[str, Any]): + obs, reward, done, truncated, info = super().step(action) + self._ensure_initialized() + if self.taxim_last_render + (1 / self.taxim_fps) <= self.data.time: + self.last_tactile_frames = self._render_tactile_frames(visualize=self.visualize) + self.taxim_last_render = self.data.time + + self._update_obs(obs) + return obs, reward, done, truncated, info diff --git a/pyproject.toml b/pyproject.toml index ed4da7fc..e472ace8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -225,6 +225,11 @@ version_files = [ "extensions/rcs_tacto/src/rcs_tacto/__init__.py:__version__", "extensions/rcs_tacto/pyproject.toml:\"rcs-core>=(.*)\"", + "extensions/rcs_taxim/pyproject.toml:version", + "extensions/rcs_taxim/src/rcs_taxim/__init__.py:__version__", + "extensions/rcs_taxim/pyproject.toml:\"rcs-core>=(.*)\"", + + "extensions/rcs_robotiq2f85/pyproject.toml:version", "extensions/rcs_robotiq2f85/src/rcs_robotiq2f85/__init__.py:__version__", "extensions/rcs_robotiq2f85/pyproject.toml:\"rcs-core>=(.*)\"", diff --git a/python/rcs/envs/configs.py b/python/rcs/envs/configs.py index 2139c7d0..6b87fdd9 100644 --- a/python/rcs/envs/configs.py +++ b/python/rcs/envs/configs.py @@ -31,7 +31,7 @@ class EmptyWorldFR3(SimEnvCreator): - robot_prefix_template = "right" + robot_prefix_template = "robot" gripper_prefix_template = "gripper" def config(self) -> SimEnvCreatorConfig: