feat(sim): replay trajectories from dynamic joint state

Makefile

@@ -32,9 +32,11 @@ stubgen:
 	find ./python -not -path "./python/rcs/_core/*" -name '*.pyi' -delete
 	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/tuple\[typing\.Literal\[\([0-9]\+\)\], typing\.Literal\[1\]\]/tuple\[typing\.Literal[\1]\]/g'
 	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/tuple\[\([M|N]\), typing\.Literal\[1\]\]/tuple\[\1\]/g'
-	sed -i 's/    q_home: numpy\.ndarray\[tuple\[M\], numpy\.dtype\[numpy\.float64\]\] | None/    q_home: numpy.ndarray | None/' python/rcs/_core/common.pyi
-	python -c "from pathlib import Path; p=Path('python/rcs/_core/common.pyi'); t=p.read_text(); t=t.replace('numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]]', 'numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]'); p.write_text(t)"
-	python -c "from pathlib import Path; p=Path('python/rcs/_core/sim.pyi'); t=p.read_text(); t=t.replace('numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]]', 'numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]'); t=t.replace(', max_buffer_frames: int = 100', ''); p.write_text(t)"
+	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/class RobotConfig/class RobotConfig(typing.Generic[M])/g'
+	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/class SimRobotConfig(rcs._core.common.RobotConfig)/class SimRobotConfig(rcs._core.common.RobotConfig[M])/g'
+	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/class DynamicJointState/class DynamicJointState(typing.Generic[M])/g'
+	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/N = typing.TypeVar("N", bound=int)//g'
+	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/, N/, M/g'
 	python ci_scripts/generate_common_typing.py
 	ruff check --fix python/rcs/_core python/rcs/common_typing.py
 	isort python/rcs/_core python/rcs/common_typing.py

extensions/rcs_fr3/pyproject.toml

@@ -19,7 +19,7 @@ dependencies = ["rcs>=0.6.3", "frankik"]
 readme = "README.md"
 maintainers = [{ name = "Tobias Jülg", email = "tobias.juelg@utn.de" }]
 authors = [{ name = "Tobias Jülg", email = "tobias.juelg@utn.de" }]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 
 [tool.scikit-build]

extensions/rcs_panda/pyproject.toml

@@ -18,7 +18,7 @@ dependencies = ["rcs>=0.6.3"]
 readme = "README.md"
 maintainers = [{ name = "Tobias Jülg", email = "tobias.juelg@utn.de" }]
 authors = [{ name = "Tobias Jülg", email = "tobias.juelg@utn.de" }]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 
 [tool.scikit-build]

extensions/rcs_robotics_library/pyproject.toml

@@ -21,7 +21,7 @@ authors = [
   { name = "Tobias Jülg", email = "tobias.juelg@utn.de" },
   { name = "Pierre Krack", email = "pierre.krack@utn.de" },
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 
 [tool.scikit-build]

extensions/rcs_robotiq2f85/pyproject.toml

@@ -17,5 +17,5 @@ maintainers = [
 authors = [
   { name = "Tobias Jülg", email = "tobias.juelg@utn.de" },
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 license = { text = "AGPL-3.0-or-later" }

extensions/rcs_so101/pyproject.toml

@@ -18,7 +18,7 @@ dependencies = ["rcs>=0.6.3", "lerobot==0.3.3"]
 readme = "README.md"
 maintainers = [{ name = "Tobias Jülg", email = "tobias.juelg@utn.de" }]
 authors = [{ name = "Tobias Jülg", email = "tobias.juelg@utn.de" }]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 [tool.scikit-build]
 build.verbose = true

extensions/rcs_tacto/pyproject.toml

@@ -17,7 +17,7 @@ maintainers = [
   { name = "Seongjin Bien", email = "seongjin.bien@utn.de" },
 ]
 authors = [{ name = "Seongjin Bien", email = "seongjin.bien@utn.de" }]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 [tool.black]
 line-length = 120

extensions/rcs_ur5e/pyproject.toml

@@ -15,7 +15,7 @@ authors = [
   { name = "Tobias Jülg", email = "tobias.juelg@utn.de" },
   { name = "Johannes Hechtl", email = "johannes.hechtl@siemens.com" },
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 [tool.black]
 line-length = 120

extensions/rcs_usb_cam/pyproject.toml

@@ -12,7 +12,7 @@ maintainers = [
   { name = "Seongjin Bien", email = "seongjin.bien@utn.de" },
 ]
 authors = [{ name = "Seongjin Bien", email = "seongjin.bien@utn.de" }]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 [tool.black]
 line-length = 120

extensions/rcs_xarm7/pyproject.toml

@@ -16,7 +16,7 @@ authors = [
   { name = "Tobias Jülg", email = "tobias.juelg@utn.de" },
   { name = "Ken Nakahara", email = "knakahara@lasr.org" },
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 [tool.black]
 line-length = 120

pyproject.toml

@@ -50,7 +50,7 @@ authors = [
   { name = "Pierre Krack", email = "pierre.krack@utn.de" },
   { name = "Seongjin Bien", email = "seongjin.bien@utn.de" },
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 license = { file = "LICENSE" }
 
 [dependency-groups]

python/rcs/_core/common.pyi

@@ -41,7 +41,6 @@ __all__: list[str] = [
     "TRIPOD_GRASP",
 ]
 M = typing.TypeVar("M", bound=int)
-N = typing.TypeVar("N", bound=int)
 
 class BaseCameraConfig:
     frame_rate: int
@@ -239,20 +238,20 @@ class Robot:
     def to_pose_in_robot_coordinates(self, pose_in_world_coordinates: Pose) -> Pose: ...
     def to_pose_in_world_coordinates(self, pose_in_robot_coordinates: Pose) -> Pose: ...
 
-class RobotConfig:
+class RobotConfig(typing.Generic[M]):
     attachment_site: str
     dof: int
-    joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]
+    joint_limits: numpy.ndarray[tuple[typing.Literal[2], M], numpy.dtype[numpy.float64]]
     kinematic_model_path: str
-    q_home: numpy.ndarray | None
+    q_home: numpy.ndarray[tuple[M], numpy.dtype[numpy.float64]] | None
     robot_platform: RobotPlatform
     robot_type: RobotType
     tcp_offset: Pose
     def __init__(
         self,
         robot_type: RobotType = ...,
         dof: int = 7,
-        joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]] = ...,
+        joint_limits: numpy.ndarray[tuple[typing.Literal[2], M], numpy.dtype[numpy.float64]] = ...,
         robot_platform: RobotPlatform = ...,
         tcp_offset: Pose = ...,
         attachment_site: str = "attachment_site",

python/rcs/_core/sim.pyi

@@ -11,6 +11,8 @@ import rcs._core.common
 
 __all__: list[str] = [
     "CameraType",
+    "DynamicJointSchema",
+    "DynamicJointState",
     "FrameSet",
     "GuiClient",
     "Sim",
@@ -32,7 +34,6 @@ __all__: list[str] = [
     "tracking",
 ]
 M = typing.TypeVar("M", bound=int)
-N = typing.TypeVar("N", bound=int)
 
 class CameraType:
     """
@@ -69,6 +70,18 @@ class CameraType:
     @property
     def value(self) -> int: ...
 
+class DynamicJointSchema:
+    joint_names: list[str]
+    joint_types: list[int]
+    qpos_sizes: list[int]
+    qvel_sizes: list[int]
+    def __init__(self) -> None: ...
+
+class DynamicJointState(typing.Generic[M]):
+    qpos: numpy.ndarray[tuple[M], numpy.dtype[numpy.float64]]
+    qvel: numpy.ndarray[tuple[M], numpy.dtype[numpy.float64]]
+    def __init__(self) -> None: ...
+
 class FrameSet:
     def __init__(
         self,
@@ -94,9 +107,12 @@ class Sim:
     def _start_gui_server(self, id: str) -> None: ...
     def _stop_gui_server(self) -> None: ...
     def get_config(self) -> SimConfig: ...
+    def get_dynamic_joint_schema(self) -> DynamicJointSchema: ...
+    def get_dynamic_joint_state(self) -> DynamicJointState: ...
     def is_converged(self) -> bool: ...
     def reset(self) -> None: ...
     def set_config(self, cfg: SimConfig) -> bool: ...
+    def set_dynamic_joint_state(self, schema: DynamicJointSchema, state: DynamicJointState) -> None: ...
     def step(self, k: int) -> None: ...
     def step_until_convergence(self) -> None: ...
     def sync_gui(self) -> None: ...
@@ -110,7 +126,9 @@ class SimCameraConfig(rcs._core.common.BaseCameraConfig):
     ) -> None: ...
 
 class SimCameraSet:
-    def __init__(self, sim: Sim, cameras: dict[str, SimCameraConfig], render_on_demand: bool = True) -> None: ...
+    def __init__(
+        self, sim: Sim, cameras: dict[str, SimCameraConfig], render_on_demand: bool = True, max_buffer_frames: int = 100
+    ) -> None: ...
     def buffer_size(self) -> int: ...
     def clear_buffer(self) -> None: ...
     def get_latest_frameset(self) -> FrameSet | None: ...
@@ -195,12 +213,12 @@ class SimRobot(rcs._core.common.Robot):
     def set_config(self, cfg: SimRobotConfig) -> bool: ...
     def set_joints_hard(self, q: numpy.ndarray[tuple[M], numpy.dtype[numpy.float64]]) -> None: ...
 
-class SimRobotConfig(rcs._core.common.RobotConfig):
+class SimRobotConfig(rcs._core.common.RobotConfig[M]):
     actuators: list[str]
     arm_collision_geoms: list[str]
     base: str
     dof: int
-    joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]
+    joint_limits: numpy.ndarray[tuple[typing.Literal[2], M], numpy.dtype[numpy.float64]]
     joint_rotational_tolerance: float
     joints: list[str]
     seconds_between_callbacks: float
@@ -247,7 +265,7 @@ class SimRobotConfig(rcs._core.common.RobotConfig):
         ],
         base: str = "base",
         dof: int = 7,
-        joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]] = ...,
+        joint_limits: numpy.ndarray[tuple[typing.Literal[2], M], numpy.dtype[numpy.float64]] = ...,
     ) -> None: ...
     def add_prefix(self, id: str) -> None: ...
 

python/rcs/envs/base.py

@@ -170,6 +170,7 @@ class ArmObsType(TQuatDictType, JointsDictType, TRPYDictType): ...
 
 CartOrJointContType: TypeAlias = TQuatDictType | JointsDictType | TRPYDictType
 LimitedCartOrJointContType: TypeAlias = LimitedTQuatRelDictType | LimitedJointsRelDictType | LimitedTRPYRelDictType
+SimStateSchema: TypeAlias = dict[str, list[str] | list[int]]
 
 
 class ArmWithGripper(TQuatDictType, GripperDictType): ...
@@ -204,18 +205,18 @@ class HardwareEnv(BaseEnv):
 class SimEnv(BaseEnv):
     PLATFORM = RobotPlatform.SIMULATION
     STATE_KEY = "sim_state"
-    STATE_SPEC_KEY = "sim_state_spec"
+    STATE_SCHEMA_KEY = "sim_state_schema"
 
     def __init__(self, sim: simulation.Sim, return_state=True) -> None:
         self.sim = sim
         cfg = self.sim.get_config()
         self.frame_rate = SimpleFrameRate(cfg.frequency, "MoJoCo Simulation Loop")
         self.main_greenlet: greenlet | None = None
         self.return_state = return_state
-        self._replay_state: tuple[np.ndarray, int | None] | None = None
+        self._replay_state: tuple[np.ndarray, SimStateSchema | None] | None = None
 
-    def set_replay_state(self, state: np.ndarray, spec: int | None = None):
-        self._replay_state = (state, spec)
+    def set_replay_state(self, state: np.ndarray, schema: SimStateSchema | None = None):
+        self._replay_state = (state, schema)
 
     def step(self, action: dict[str, Any]) -> tuple[dict[str, Any], float, bool, bool, dict]:
         if self.main_greenlet is not None:
@@ -255,7 +256,7 @@ def reset(
     def observation(self, observation: dict[str, Any], info: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
         sim_state = self.sim.get_state()
         info[self.STATE_KEY] = sim_state
-        info[self.STATE_SPEC_KEY] = self.sim.get_state_spec()
+        info[self.STATE_SCHEMA_KEY] = self.sim.get_state_schema()
         return observation, info
 
 

python/rcs/envs/configs.py

@@ -1,6 +1,6 @@
 import copy
 import time
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 import numpy as np
 from rcs._core.common import FrankaHandTCPOffset, GripperType, RobotType
@@ -37,7 +37,7 @@ class EmptyWorldFR3(SimEnvCreator):
     def config(self) -> SimEnvCreatorConfig:
         q_home = rcs.ROBOTS[RobotType.FR3].q_home
         q_home[-1] = np.pi / 4
-        robot_cfg = SimRobotConfig(
+        robot_cfg: SimRobotConfig[Literal[7]] = SimRobotConfig(
             robot_type=RobotType.FR3,
             tcp_offset=GRIPPER_OFFSETS[rcs.common.GripperType.FrankaHand],
             attachment_site=rcs.ROBOTS[RobotType.FR3].attachment_site,
@@ -183,7 +183,7 @@ class EmptyWorldFR3Duo(SimEnvCreator):
     gripper_mesh_quaternion_offset: ClassVar[list[float]] = [0, 0, 0.7071068, 0.7071068]
 
     def config(self) -> SimEnvCreatorConfig:
-        robot_cfg = SimRobotConfig(
+        robot_cfg: SimRobotConfig[Literal[7]] = SimRobotConfig(
             tcp_offset=GRIPPER_OFFSETS[rcs.common.GripperType("Robotiq2F85")],
             robot_type=RobotType.FR3,
             attachment_site=rcs.ROBOTS[RobotType.FR3].attachment_site,
@@ -224,7 +224,7 @@ def config(self) -> SimEnvCreatorConfig:
             joint_limits=rcs.ROBOTS[RobotType.FR3].joint_limits,
             q_home=rcs.HOME_POSITIONS["FR3_DUO_LEFT"],
         )
-        robot_cfg_right = copy.deepcopy(robot_cfg)
+        robot_cfg_right: SimRobotConfig[Literal[7]] = copy.deepcopy(robot_cfg)
         robot_cfg_right.q_home = rcs.HOME_POSITIONS["FR3_DUO_RIGHT"]
 
         robot_cfgs: dict[str, SimRobotConfig] = {"left": robot_cfg, "right": robot_cfg_right}

python/rcs/envs/scenes.py

@@ -49,6 +49,7 @@ def __call__(self, **kwargs) -> gym.Env:
 class WrapperConfig:
     binary_gripper: bool = True
     home_on_reset: bool = True
+    include_depth: bool = False
 
 
 #### SIM SPECIFIC ####
@@ -239,7 +240,13 @@ def create_model(self, cfg: SimEnvCreatorConfig) -> MjModel:
         if cfg.root_frame_objects is not None:
             for object_id, (object_xml, object2root_frame) in cfg.root_frame_objects.items():
                 object2world = cfg.root_frame_to_world * object2root_frame
-                self.add_object_mujoco(composer, object_id, object_xml, object2world)
+                self.add_object_mujoco(
+                    composer,
+                    object_id,
+                    object_xml,
+                    object2world,
+                    register_root_relative_replay_free_joints=True,
+                )
         # add external objects
         if cfg.world_frame_objects is not None:
             for object_id, (object_xml, object2world) in cfg.world_frame_objects.items():
@@ -324,6 +331,11 @@ def create_env_from_model(self, cfg: SimEnvCreatorConfig, mjmodel: MjModel) -> g
             prefixed_cfg = self.prefixed_cfg(cfg)
 
         simulation = Sim(mjmodel, prefixed_cfg.sim_cfg)
+        if isinstance(mjmodel, ModelComposer):
+            simulation.configure_state_encodings(
+                root_frame_to_world=cfg.root_frame_to_world,
+                root_relative_free_joints=mjmodel.root_relative_replay_free_joints,
+            )
 
         envs: dict[str, gym.Env] = {}
         env: gym.Env
@@ -353,7 +365,7 @@ def create_env_from_model(self, cfg: SimEnvCreatorConfig, mjmodel: MjModel) -> g
                 BaseCameraSet,
                 SimCameraSet(simulation, prefixed_cfg.camera_cfgs, physical_units=True, render_on_demand=True),
             )
-            env = CameraSetWrapper(env, camera_set, include_depth=True)
+            env = CameraSetWrapper(env, camera_set, include_depth=cfg.wrapper_cfg.include_depth)
         env = self.add_task_env(prefixed_cfg.task_cfg, env, simulation, cfg)
         if not prefixed_cfg.headless:
             env.get_wrapper_attr("sim").open_gui()
@@ -373,13 +385,20 @@ def add_task_env(
         return env
 
     def add_object_mujoco(
-        self, composer: ModelComposer, object_id: str, object_xml: str, object2world: rcs.common.Pose
+        self,
+        composer: ModelComposer,
+        object_id: str,
+        object_xml: str,
+        object2world: rcs.common.Pose,
+        *,
+        register_root_relative_replay_free_joints: bool = False,
     ):
         """Add an object to the Mujoco scene."""
         composer.add_object_world_frame(
             object_xml,
             object_prefix=object_id + "_",
             pose=object2world,
+            register_root_relative_replay_free_joints=register_root_relative_replay_free_joints,
         )
 
     def add_object_robot_frame_mujoco(

python/rcs/envs/tasks.py

@@ -167,6 +167,7 @@ def add_task_mujoco(cfg: PickTaskConfig, composer: ModelComposer, env_cfg: SimEn
             cfg.object_xml,
             object_prefix=cfg.prefix,
             pose=object2world,
+            register_root_relative_replay_free_joints=True,
         )
 
     @staticmethod