Add articulator to the recording system

README.md

@@ -46,7 +46,7 @@ factors.
          - Frames appear synced, however when compared to real-world footage there is a speed scaling issue...
     - [x] Validate "Replay episode" works as expected with episodes collected in Isaac Sim
     - [ ] Better calibrate robots so they match position in sim to real world
-    - [ ] Record position of articulator in sim, not just real robot joints
+    - [x] Record position of articulator in sim, not just real robot joints
     - [x] Fix bug with myarm firmware where there's a singularity at the 0 point
     - [ ] Investigate high latency in myarm loop (it is higher than in lerobot branch)
   - [x] Create a RobotProtocol that emulates latency and speed of my real robot

isaac_src/extensions/trajectory_synth/exts/trajectory_synth/trajectory_synth/lerobot_episode.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Literal
 
 from . import schemas
 
@@ -40,8 +41,18 @@ def add_timestep(self, isaac_time: float):
         )
 
         # Interpolate joints
-        leader_joints = self._interpolate_joints(isaac_time, self.leader_steps)
-        follower_joints = self._interpolate_joints(isaac_time, self.follower_steps)
+        leader_ros_joints = self._interpolate_joints(
+            isaac_time, self.leader_steps, "real"
+        )
+        leader_sim_joints = self._interpolate_joints(
+            isaac_time, self.leader_steps, "sim"
+        )
+        follower_ros_joints = self._interpolate_joints(
+            isaac_time, self.follower_steps, "real"
+        )
+        follower_sim_joints = self._interpolate_joints(
+            isaac_time, self.follower_steps, "sim"
+        )
 
         # Modify isaac_time so it stretches or squishes such that 'target_fps' is hit
         # TODO: I need to look into this further to validate it's working as intended
@@ -50,8 +61,10 @@ def add_timestep(self, isaac_time: float):
         # Create a LeRobotTimestep
         # (Here we store leader joints in .action, follower in .state)
         new_timestep = schemas.LeRobotTimestep(
-            action=leader_joints,
-            state=follower_joints,
+            ros_action=leader_ros_joints,
+            ros_state=follower_ros_joints,
+            sim_action=leader_sim_joints,
+            sim_state=follower_sim_joints,
             timestamp=timestamp,
             frame_index=frame_index,
             episode_index=self.episode_index,
@@ -62,7 +75,10 @@ def add_timestep(self, isaac_time: float):
         self.timesteps.append(new_timestep)
 
     def _interpolate_joints(
-        self, time: float, time_samples: list[schemas.RobotTimeSample]
+        self,
+        time: float,
+        time_samples: list[schemas.RobotTimeSample],
+        joints_source: Literal["real", "sim"],
     ) -> list[float]:
         """
         Given a time and a particular robot_key (e.g. 'leader' or 'follower'),
@@ -86,7 +102,7 @@ def _interpolate_joints(
 
         # Validate time is in-range
         if time <= time_samples[0].isaac_time:
-            return time_samples[0].joint_positions
+            return time_samples[0].from_source(joints_source)
         if time >= time_samples[-1].isaac_time:
             raise ValueError(f"Time {time} is after the last sample")
 
@@ -98,16 +114,16 @@ def _interpolate_joints(
                 # Interpolate linearly
                 ratio = (time - t0) / (t1 - t0)
 
-                joints0 = time_samples[i].joint_positions
-                joints1 = time_samples[i + 1].joint_positions
+                joints0 = time_samples[i].from_source(joints_source)
+                joints1 = time_samples[i + 1].from_source(joints_source)
                 interp_joints = [
                     j0 + ratio * (j1 - j0)
                     for j0, j1 in zip(joints0, joints1, strict=True)
                 ]
                 return interp_joints
 
         # Fallback (should not happen if times cover the entire range)
-        return time_samples[-1].joint_positions
+        return time_samples[-1].from_source(joints_source)
 
     def save(self):
         """Save the list as json lines"""

isaac_src/extensions/trajectory_synth/exts/trajectory_synth/trajectory_synth/recorder.py

@@ -7,6 +7,7 @@
 import omni.graph.core.types as ot
 import omni.kit.commands
 from omni import ui
+from omni.isaac.core.articulations import Articulation
 
 from . import path_utils, schemas
 from .episode import EpisodeRecording
@@ -226,6 +227,10 @@ def set_up_omnigraph(self, robot_attributes: schemas.RobotAttributes):
         # so we can reinitialize the extension without conflicts
         recorder_name = f"JointRecorder_{random.randint(0, 10000000)}"
 
+        # The custom node needs a reference to the sim robot, so we can record not just
+        # the real robots joints, but also the simulated robot's joints
+        sim_articulator = Articulation(prim_path=robot_attributes.root_joint)
+
         @og.create_node_type(unique_name=recorder_name)
         def custom_joint_state_processor(
             input_joint_names: ot.tokenarray,
@@ -244,11 +249,16 @@ def custom_joint_state_processor(
                 logging.debug("No joint positions received from ROS yet.")
                 return ""
 
+            if not sim_articulator.handles_initialized:
+                sim_articulator.initialize()
+                logging.info(f"Initialized sim articulator for {robot_attributes.name}")
+
             robot_recording = self.current_joint_recording.robots[robot_attributes.name]
             robot_recording.joint_names = input_joint_names
             robot_recording.time_samples.append(
                 schemas.RobotTimeSample(
-                    joint_positions=input_positions,
+                    real_joint_positions=input_positions,
+                    sim_joint_positions=sim_articulator.get_joint_positions().tolist(),
                     isaac_time=timestamp,
                     ros_time=ros_timestamp,
                 )

isaac_src/extensions/trajectory_synth/exts/trajectory_synth/trajectory_synth/schemas/episode_files.py

@@ -1,13 +1,21 @@
+from typing import Literal
+
 from pydantic import BaseModel
 
 
 class RobotTimeSample(BaseModel):
     """A recording for a single robot, for a single episode"""
 
-    joint_positions: list[float]
+    real_joint_positions: list[float]
+    sim_joint_positions: list[float]
     isaac_time: float
     ros_time: float
 
+    def from_source(self, source: Literal["real", "sim"]) -> list[float]:
+        return (
+            self.real_joint_positions if source == "real" else self.sim_joint_positions
+        )
+
 
 class RobotRecording(BaseModel):
     """A recording for a single robot, for a single episode"""
@@ -34,8 +42,10 @@ class TrajectoryRecordingMeta(BaseModel):
 class LeRobotTimestep(BaseModel):
     """Closely mimics one row of a episode *.parquet file from a lerobot dataset"""
 
-    action: list[float]
-    state: list[float]
+    ros_action: list[float]
+    sim_action: list[float]
+    ros_state: list[float]
+    sim_state: list[float]
     timestamp: float
     frame_index: int
     episode_index: int

pkgs/myarm_ai/myarm_ai/datasets/traj_synth/reader.py

@@ -80,10 +80,11 @@ def timesteps(self) -> Generator[Timestep, None, None]:
                 obs_dir.name: self._image_for_observation(step["frame_index"], obs_dir)
                 for obs_dir in self._observation_paths
             }
+            # TODO: Whether joints are pulled from 'ros' or 'sim' should be configurable
             yield Timestep(
                 images=images,
-                action=step["action"],
-                state=step["state"],  # type: ignore
+                action=step["ros_action"],
+                state=step["ros_state"],  # type: ignore
                 timestamp=step["timestamp"],
             )