feat: add lerobot rcs eval config

Author: juelg

python/rcs/__init__.py

@@ -2,17 +2,19 @@
 
 import os
 import site
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from gymnasium import register
+from lerobot.configs.types import PolicyFeature
+from lerobot.envs.configs import EnvConfig
 from rcs._core import __version__, common
 from rcs.envs.creators import (
     FR3LabDigitGripperPickUpSimEnvCreator,
     FR3SimplePickUpSimEnvCreator,
 )
+from rcs.lerobot import Pick
 
 from rcs import camera, envs, hand, sim
-from rcs.lerobot import Pick
 
 
 @dataclass(kw_only=True)
@@ -105,11 +107,36 @@ def get_scene_urdf(scene_name: str) -> str | None:
     entry_point=Pick(),
 )
 
-# Genius TODO: Add the tacto version of the SimEnvCreator
 # TODO: gym.make("rcs/FR3SimEnv-v0") results in a pickling error:
 # TypeError: cannot pickle 'rcs._core.sim.SimRobotConfig' object
 # cf. https://pybind11.readthedocs.io/en/stable/advanced/classes.html#deepcopy-support
 # register(
 #    id="rcs/FR3SimEnv-v0",
 #    entry_point=SimEnvCreator(),
 # )
+
+
+# lerobot env registration
+@EnvConfig.register_subclass("rcs/pick")
+@dataclass
+class RCSEnvConfig(EnvConfig):
+    task: str | None = "pick"
+    fps: int = 30
+    features: dict[str, PolicyFeature] = field(default_factory=dict)
+    features_map: dict[str, str] = field(default_factory=dict)
+    max_parallel_tasks: int = 1
+    disable_env_checker: bool = True
+
+    @property
+    def package_name(self) -> str:
+        """Package name to import if environment not found in gym registry"""
+        return "rcs"
+
+    @property
+    def gym_id(self) -> str:
+        """ID string used in gym.make() to instantiate the environment"""
+        return f"{self.package_name}/{self.task}"
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {}

python/rcs/lerobot.py

@@ -1,18 +1,16 @@
 from typing import Any
+
 import gymnasium as gym
-import rcs
 import numpy as np
-
-from rcs._core.sim import SimConfig
+from gymnasium.envs.registration import EnvCreator
+from rcs._core.sim import CameraType, SimCameraConfig, SimConfig
 from rcs.envs.base import ControlMode, RelativeActionSpace
+from rcs.envs.creators import SimTaskEnvCreator
 from rcs.envs.space_utils import ActObsInfoWrapper
 from rcs.sim import SimGripperConfig
 from rcs_tacto.tacto_wrapper import TactoSimWrapper
-from gymnasium.envs.registration import EnvCreator
 
-
-from rcs._core.sim import CameraType, SimCameraConfig
-from rcs.envs.creators import SimTaskEnvCreator
+import rcs
 
 SCENE_FILE = "beast_refiner/rcs_env/rcs_icra_scene_digit/scene.xml"
 CAMERAS = [
@@ -23,17 +21,20 @@
     "side_wide",
 ]
 
+
 class AlignSpaceWrapper(ActObsInfoWrapper):
 
     def __init__(self, env):
         super().__init__(env)
-        self.observation_space = gym.spaces.Dict({
-            "state": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(15,), dtype=np.float32),
-            "images.image": env.observation_space["frames"]["side"]["rgb"]["data"],
-            "images.image2": env.observation_space["frames"]["wrist"]["rgb"]["data"],
-            "images.tactile_left": gym.spaces.Box(low=0, high=255, shape=(320, 240, 3), dtype=np.uint8),
-            "images.tactile_right": gym.spaces.Box(low=0, high=255, shape=(320, 240, 3), dtype=np.uint8),
-        })
+        self.observation_space = gym.spaces.Dict(
+            {
+                "state": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(15,), dtype=np.float32),
+                "images.image": env.observation_space["frames"]["side"]["rgb"]["data"],
+                "images.image2": env.observation_space["frames"]["wrist"]["rgb"]["data"],
+                "images.tactile_left": gym.spaces.Box(low=0, high=255, shape=(320, 240, 3), dtype=np.uint8),
+                "images.tactile_right": gym.spaces.Box(low=0, high=255, shape=(320, 240, 3), dtype=np.uint8),
+            }
+        )
         # self.action_space = env.action_space["joints"]
         gym.spaces.Box(low=-np.inf, high=np.inf, shape=(8,), dtype=np.float32)
 
@@ -48,7 +49,6 @@ def observation(self, observation: dict[str, Any], info: dict[str, Any]) -> tupl
         # tactile and tau_ext are still unsolved
 
         return observation, info
-        
 
     def action(self, action: np.ndarray) -> dict[str, Any]:
         action = {"joints": action[:7], "gripper": action[7]}
@@ -151,4 +151,4 @@ def make_env(n_envs: int = 1, use_async_envs: bool = False):
 
 class Pick(EnvCreator):
     def __call__(self) -> gym.Env:
-        return make_single_env()
+        return make_single_env()