Merge pull request #156 from utn-mi/juelg/register-gym-envs

Register envs in gym

Author: juelg

CMakeLists.txt

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

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = ["websockets>=11.0",
     "python-dotenv==1.0.1",
     "opencv-python~=4.10.0.84",
     # NOTE: when changing the mujoco version, also change it in requirements_dev.txt
-    "mujoco==3.1.5"
+    "mujoco==3.2.6"
 ]
 readme = "README.md"
 maintainers = [
@@ -98,6 +98,7 @@ ignore = [
   "T201",    # print() used
   "PLR5501", # elif to reduce indentation
   "PT018",   # assertion should be broken down into multiple parts
+  "NPY002",
 ]
 
 [tool.pylint.format]

python/examples/fr3.py

@@ -79,13 +79,8 @@ def main():
 
             # add camera to have a rendering gui
             cameras = {
-                "default_free": SimCameraConfig(
-                    identifier="", type=int(CameraType.default_free)
-                ),
-                "wrist": SimCameraConfig(
-                    identifier="eye-in-hand_0", type=int(CameraType.fixed)
-                ),
-                # TODO: odd behavior when not both cameras are used: only last image is shown
+                "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free)),
+                "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),
             }
             cam_cfg = SimCameraSetConfig(cameras=cameras, resolution_width=1280, resolution_height=720, frame_rate=20)
             camera_set = SimCameraSet(simulation, cam_cfg)

python/examples/grasp_demo.py

@@ -0,0 +1,93 @@
+import logging
+
+import gymnasium as gym
+import mujoco
+import numpy as np
+from rcsss._core.common import Pose
+from rcsss._core.sim import FR3State
+from rcsss.envs.base import GripperWrapper
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+class PickUpDemo:
+    def __init__(self, env: gym.Env):
+        self.env = env
+        self.home_pose = self.env.unwrapped.robot.get_cartesian_position()
+
+    def _action(self, pose: Pose, gripper: float) -> np.ndarray:
+        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)
+        return Pose(translation=data.geom_xpos[geom_id], rotation=data.geom_xmat[geom_id].reshape(3, 3))
+
+    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:
+        re = self.env.step(action)
+        s: FR3State = self.env.unwrapped.robot.get_state()
+        return re[0]
+
+    def plan_linear_motion(self, geom_name: str, delta_up: float, num_waypoints: int = 20) -> list[Pose]:
+        end_eff_pose = self.env.unwrapped.robot.get_cartesian_position()
+
+        goal_pose = self.get_object_pose(geom_name=geom_name)
+        # goal pose is above the object and gripper coordinate must flip z-axis (end effector base rotation is [1, 0, 0, 0])
+        # be careful we define identity quaternion as as [0, 0, 0, 1]
+        # this does not work if the object is flipped
+        goal_pose *= Pose(translation=[0, 0, delta_up], quaternion=[1, 0, 0, 0])
+
+        waypoints = self.generate_waypoints(end_eff_pose, goal_pose, num_waypoints=num_waypoints)
+        return 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()
+            obs = self.step(self._action(delta_action, 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)
+        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=15)
+        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=20)
+        self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
+
+    def move_home(self):
+        end_eff_pose = self.env.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/SimplePickUpSim-v0", render_mode="human", delta_actions=True)
+    env.reset()
+    controller = PickUpDemo(env)
+    controller.pickup("yellow_box_geom")
+
+
+if __name__ == "__main__":
+    main()

python/rcsss/__init__.py

@@ -3,11 +3,25 @@
 import pathlib
 import site
 
+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
 
+# available mujoco scenes
 scenes = {
     path.stem: path / "scene.mjb" for path in (pathlib.Path(site.getsitepackages()[0]) / "rcsss" / "scenes").glob("*")
 }
 
+# make submodules available
 __all__ = ["__doc__", "__version__", "common", "hw", "sim", "camera", "scenes", "control", "envs"]
+
+# register gymnasium environments
+register(
+    id="rcs/SimplePickUpSim-v0",
+    entry_point=FR3SimplePickUpSim(),
+)
+register(
+    id="rcs/FR3Real-v0",
+    entry_point=FR3Real(),
+)

python/rcsss/envs/base.py

@@ -523,6 +523,10 @@ def close(self):
 
 class GripperWrapper(ActObsInfoWrapper):
     # TODO: sticky gripper, like in aloha
+
+    BINARY_GRIPPER_CLOSED = 0
+    BINARY_GRIPPER_OPEN = 1
+
     def __init__(self, env, gripper: common.Gripper, binary: bool = True):
         super().__init__(env)
         self.unwrapped: FR3Env
@@ -543,7 +547,9 @@ def reset(self, **kwargs) -> tuple[dict[str, Any], dict[str, Any]]:
     def observation(self, observation: dict[str, Any], info: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
         observation = copy.deepcopy(observation)
         if self.binary:
-            observation[self.gripper_key] = self._last_gripper_cmd if self._last_gripper_cmd is not None else 1
+            observation[self.gripper_key] = (
+                self._last_gripper_cmd if self._last_gripper_cmd is not None else self.BINARY_GRIPPER_OPEN
+            )
         else:
             observation[self.gripper_key] = self._gripper.get_normalized_width()
 
@@ -565,7 +571,7 @@ def action(self, action: dict[str, Any]) -> dict[str, Any]:
 
         if self._last_gripper_cmd is None or self._last_gripper_cmd != gripper_action:
             if self.binary:
-                self._gripper.grasp() if gripper_action == 0 else self._gripper.open()
+                self._gripper.grasp() if gripper_action == self.BINARY_GRIPPER_CLOSED else self._gripper.open()
             else:
                 self._gripper.set_normalized_width(gripper_action)
         self._last_gripper_cmd = gripper_action

python/rcsss/envs/factories.py

@@ -1,8 +1,11 @@
 import logging
 from os import PathLike
+from typing import Type
 
 import gymnasium as gym
+import numpy as np
 import rcsss
+from gymnasium.envs.registration import EnvCreator
 from rcsss import sim
 from rcsss._core.hw import FR3Config, IKController
 from rcsss._core.sim import CameraType
@@ -21,7 +24,13 @@
     RelativeTo,
 )
 from rcsss.envs.hw import FR3HW
-from rcsss.envs.sim import CollisionGuard, FR3Sim
+from rcsss.envs.sim import (
+    CollisionGuard,
+    FR3Sim,
+    FR3SimplePickUpSimSuccessWrapper,
+    RandomCubePos,
+    SimWrapper,
+)
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -43,7 +52,9 @@ def default_fr3_hw_gripper_cfg():
     return gripper_cfg
 
 
-def default_realsense(name2id: dict[str, str]):
+def default_realsense(name2id: dict[str, str] | None) -> RealSenseCameraSet | None:
+    if name2id is None:
+        return None
     cameras = {name: BaseCameraConfig(identifier=id) for name, id in name2id.items()}
     cam_cfg = RealSenseSetConfig(
         cameras=cameras,
@@ -132,6 +143,7 @@ def fr3_hw_env(
             control_mode=control_mode,
             tcp_offset=rcsss.common.Pose(rcsss.common.FrankaHandTCPOffset()),
             sim_gui=True,
+            truncate_on_collision=False,
         )
     if max_relative_movement is not None:
         env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
@@ -173,6 +185,7 @@ def fr3_sim_env(
     relative_to: RelativeTo = RelativeTo.LAST_STEP,
     urdf_path: str | PathLike | None = None,
     mjcf: str | PathLike = "fr3_empty_world",
+    sim_wrapper: Type[SimWrapper] | None = None,
 ) -> gym.Env[ObsArmsGrCam, LimitedJointsRelDictType]:
     """
     Creates a simulation environment for the FR3 robot.
@@ -190,6 +203,8 @@ def fr3_sim_env(
         urdf_path (str | PathLike | None): Path to the URDF file. If None, the URDF file is automatically deduced
             which requires a UTN compatible lab scene to be present.
         mjcf (str | PathLike): Path to the Mujoco scene XML file. Defaults to "fr3_empty_world".
+        sim_wrapper (Type[SimWrapper] | None): Wrapper to be applied before the simulation wrapper. This is useful
+            for reset management e.g. resetting objects in the scene with correct observations. Defaults to None.
 
     Returns:
         gym.Env: The configured simulation environment for the FR3 robot.
@@ -204,15 +219,15 @@ def fr3_sim_env(
     robot = rcsss.sim.FR3(simulation, "0", ik)
     robot.set_parameters(robot_cfg)
     env: gym.Env = FR3Env(robot, control_mode)
-    env = FR3Sim(env, simulation)
+    env = FR3Sim(env, simulation, sim_wrapper)
 
     if camera_set_cfg is not None:
         camera_set = SimCameraSet(simulation, camera_set_cfg)
         env = CameraSetWrapper(env, camera_set, include_depth=True)
 
     if gripper_cfg is not None:
         gripper = sim.FrankaHand(simulation, "0", gripper_cfg)
-        env = GripperWrapper(env, gripper, binary=False)
+        env = GripperWrapper(env, gripper, binary=True)
 
     if collision_guard:
         env = CollisionGuard.env_from_xml_paths(
@@ -224,8 +239,85 @@ def fr3_sim_env(
             control_mode=control_mode,
             tcp_offset=rcsss.common.Pose(rcsss.common.FrankaHandTCPOffset()),
             sim_gui=True,
+            truncate_on_collision=True,
         )
     if max_relative_movement is not None:
         env = RelativeActionSpace(env, max_mov=max_relative_movement, relative_to=relative_to)
 
     return env
+
+
+class FR3Real(EnvCreator):
+    def __call__(  # type: ignore
+        self,
+        robot_ip: str,
+        control_mode: str = "xyzrpy",
+        delta_actions: bool = True,
+        camera_config: dict[str, str] | None = None,
+        gripper: bool = True,
+    ) -> gym.Env:
+        camera_set = default_realsense(camera_config)
+        return fr3_hw_env(
+            ip=robot_ip,
+            camera_set=camera_set,
+            control_mode=(
+                ControlMode.CARTESIAN_TRPY
+                if control_mode == "xyzrpy"
+                else ControlMode.JOINTS if control_mode == "joints" else ControlMode.CARTESIAN_TQuart
+            ),
+            robot_cfg=default_fr3_hw_robot_cfg(),
+            collision_guard=None,
+            gripper_cfg=default_fr3_hw_gripper_cfg() if gripper else None,
+            max_relative_movement=(0.2, np.deg2rad(45)) if delta_actions else None,
+            relative_to=RelativeTo.LAST_STEP,
+        )
+
+
+class FR3SimplePickUpSim(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 = {
+            "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),
+            "bird_eye": SimCameraConfig(identifier="bird-eye-cam", type=int(CameraType.fixed)),
+            "side": SimCameraConfig(identifier="side_view", type=int(CameraType.fixed)),
+            "right_side": SimCameraConfig(identifier="right_side", type=int(CameraType.fixed)),
+            "left_side": SimCameraConfig(identifier="left_side", type=int(CameraType.fixed)),
+            "front": SimCameraConfig(identifier="front", 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",
+            sim_wrapper=RandomCubePos,
+        )
+        env_rel = FR3SimplePickUpSimSuccessWrapper(env_rel)
+        if render_mode == "human":
+            env_rel.get_wrapper_attr("sim").open_gui()
+
+        return env_rel