Fix joint command resampling and conversion

ddlitlab2024/dataset/converters/synced_data_converter.py

@@ -1,19 +1,7 @@
 from ddlitlab2024.dataset.converters.converter import Converter
-from ddlitlab2024.dataset.imports.data import InputData, ModelData
+from ddlitlab2024.dataset.imports.data import InputData, ModelData, joints_dict_from_msg_data
 from ddlitlab2024.dataset.models import JointCommands, JointStates, Recording, Rotation
 from ddlitlab2024.dataset.resampling.previous_interpolation_resampler import PreviousInterpolationResampler
-from ddlitlab2024.utils.utils import camelcase_to_snakecase, shift_radian_to_positive_range
-
-
-def joints_dict_from_msg_data(joints_data: list[tuple[str, float]]) -> dict[str, float]:
-    joints_dict = {}
-
-    for name, position in joints_data:
-        key = camelcase_to_snakecase(name)
-        value = shift_radian_to_positive_range(position)
-        joints_dict[key] = value
-
-    return joints_dict
 
 
 class SyncedDataConverter(Converter):
@@ -25,7 +13,9 @@ def populate_recording_metadata(self, data: InputData, recording: Recording):
 
     def convert_to_model(self, data: InputData, relative_timestamp: float, recording: Recording) -> ModelData:
         assert data.joint_state is not None, "joint_states are required in synced resampling data"
-        assert data.joint_command is not None, "joint_commands are required in synced resampling data"
+        assert all(
+            command is not None for command in data.joint_command.values()
+        ), "joint_commands are required in synced resampling data"
         assert data.rotation is not None, "IMU rotation is required in synced resampling data"
 
         models = ModelData()
@@ -50,21 +40,13 @@ def _create_rotation(self, msg, sampling_timestamp: float, recording: Recording)
         )
 
     def _create_joint_states(self, msg, sampling_timestamp: float, recording: Recording) -> JointStates:
-        if msg is None:
-            return JointStates(stamp=sampling_timestamp, recording=recording)
-        else:
-            joint_states_data = list(zip(msg.name, msg.position))
+        joint_states_data = list(zip(msg.name, msg.position))
 
-            return JointStates(
-                stamp=sampling_timestamp, recording=recording, **joints_dict_from_msg_data(joint_states_data)
-            )
-
-    def _create_joint_commands(self, msg, sampling_timestamp: float, recording: Recording) -> JointCommands:
-        if msg is None:
-            return JointCommands(stamp=sampling_timestamp, recording=recording)
-        else:
-            joint_commands_data = list(zip(msg.joint_names, msg.positions))
+        return JointStates(
+            stamp=sampling_timestamp, recording=recording, **joints_dict_from_msg_data(joint_states_data)
+        )
 
-            return JointCommands(
-                stamp=sampling_timestamp, recording=recording, **joints_dict_from_msg_data(joint_commands_data)
-            )
+    def _create_joint_commands(
+        self, joint_commands_data, sampling_timestamp: float, recording: Recording
+    ) -> JointCommands:
+        return JointCommands(stamp=sampling_timestamp, recording=recording, **joint_commands_data)

ddlitlab2024/dataset/imports/data.py

@@ -2,6 +2,18 @@
 from typing import Any
 
 from ddlitlab2024.dataset.models import GameState, Image, JointCommands, JointStates, Recording, Rotation
+from ddlitlab2024.utils.utils import camelcase_to_snakecase, shift_radian_to_positive_range
+
+
+def joints_dict_from_msg_data(joints_data: list[tuple[str, float]]) -> dict[str, float]:
+    joints_dict = {}
+
+    for name, position in joints_data:
+        key = camelcase_to_snakecase(name)
+        value = shift_radian_to_positive_range(position)
+        joints_dict[key] = value
+
+    return joints_dict
 
 
 @dataclass
@@ -18,9 +30,63 @@ class InputData:
     image: Any = None
     game_state: Any = None
     joint_state: Any = None
-    joint_command: Any = None
     rotation: Any = None
 
+    # as we are not always sending joint commands for all joints at once
+    # we need to separate them here, to enable resampling on a per joint basis
+    r_shoulder_pitch_command: Any = None
+    l_shoulder_pitch_command: Any = None
+    r_shoulder_roll_command: Any = None
+    l_shoulder_roll_command: Any = None
+    r_elbow_command: Any = None
+    l_elbow_command: Any = None
+    r_hip_yaw_command: Any = None
+    l_hip_yaw_command: Any = None
+    r_hip_roll_command: Any = None
+    l_hip_roll_command: Any = None
+    r_hip_pitch_command: Any = None
+    l_hip_pitch_command: Any = None
+    r_knee_command: Any = None
+    l_knee_command: Any = None
+    r_ankle_pitch_command: Any = None
+    l_ankle_pitch_command: Any = None
+    r_ankle_roll_command: Any = None
+    l_ankle_roll_command: Any = None
+    head_pan_command: Any = None
+    head_tilt_command: Any = None
+
+    @property
+    def joint_command(self):
+        return {
+            "r_shoulder_pitch": self.r_shoulder_pitch_command,
+            "l_shoulder_pitch": self.l_shoulder_pitch_command,
+            "r_shoulder_roll": self.r_shoulder_roll_command,
+            "l_shoulder_roll": self.l_shoulder_roll_command,
+            "r_elbow": self.r_elbow_command,
+            "l_elbow": self.l_elbow_command,
+            "r_hip_yaw": self.r_hip_yaw_command,
+            "l_hip_yaw": self.l_hip_yaw_command,
+            "r_hip_roll": self.r_hip_roll_command,
+            "l_hip_roll": self.l_hip_roll_command,
+            "r_hip_pitch": self.r_hip_pitch_command,
+            "l_hip_pitch": self.l_hip_pitch_command,
+            "r_knee": self.r_knee_command,
+            "l_knee": self.l_knee_command,
+            "r_ankle_pitch": self.r_ankle_pitch_command,
+            "l_ankle_pitch": self.l_ankle_pitch_command,
+            "r_ankle_roll": self.r_ankle_roll_command,
+            "l_ankle_roll": self.l_ankle_roll_command,
+            "head_pan": self.head_pan_command,
+            "head_tilt": self.head_tilt_command,
+        }
+
+    @joint_command.setter
+    def joint_command(self, msg):
+        joint_commands_data = list(zip(msg.joint_names, msg.positions))
+
+        for joint_name, command in joints_dict_from_msg_data(joint_commands_data).items():
+            setattr(self, f"{joint_name}_command", command)
+
 
 @dataclass
 class ModelData:

ddlitlab2024/dataset/imports/model_importer.py

@@ -31,5 +31,11 @@ def __init__(self, db: Database, strategy: ImportStrategy):
 
     def import_to_db(self, file_path: Path):
         model_data: ModelData = self.strategy.convert_to_model_data(file_path)
+
+        required_fields = ['images', 'game_states', 'joint_states', 'joint_commands']
+        for field in required_fields:
+            if not len(getattr(model_data, field)):
+                raise ValueError(f"No {field} models extracted from the file, aborting import.")
+
         self.db.session.add_all(model_data.model_instances())
         self.db.session.commit()

ddlitlab2024/dataset/imports/strategies/bitbots.py

@@ -126,7 +126,8 @@ def _create_models(self, converter: Converter, data: InputData, relative_timesta
         return self.model_data
 
     def _is_all_synced_data_available(self, data: InputData) -> bool:
-        return data.joint_command is not None and data.joint_state is not None and data.rotation is not None
+        commands_for_all_joints_available = all(command is not None for command in data.joint_command.values())
+        return commands_for_all_joints_available and data.joint_state is not None and data.rotation is not None
 
     def _create_recording(self, summary: Summary, mcap_file_path: Path) -> Recording:
         start_timestamp, end_timestamp = self._extract_timeframe(summary)

ddlitlab2024/dataset/models.py

@@ -147,26 +147,26 @@ class JointStates(Base):
     _id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
     stamp: Mapped[float] = mapped_column(Float, nullable=False)
     recording_id: Mapped[int] = mapped_column(Integer, ForeignKey("Recording._id"), nullable=False)
-    r_shoulder_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="RShoulderPitch")
-    l_shoulder_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="LShoulderPitch")
-    r_shoulder_roll: Mapped[float] = mapped_column(Float, default=0.0, name="RShoulderRoll")
-    l_shoulder_roll: Mapped[float] = mapped_column(Float, default=0.0, name="LShoulderRoll")
-    r_elbow: Mapped[float] = mapped_column(Float, default=0.0, name="RElbow")
-    l_elbow: Mapped[float] = mapped_column(Float, default=0.0, name="LElbow")
-    r_hip_yaw: Mapped[float] = mapped_column(Float, default=0.0, name="RHipYaw")
-    l_hip_yaw: Mapped[float] = mapped_column(Float, default=0.0, name="LHipYaw")
-    r_hip_roll: Mapped[float] = mapped_column(Float, default=0.0, name="RHipRoll")
-    l_hip_roll: Mapped[float] = mapped_column(Float, default=0.0, name="LHipRoll")
-    r_hip_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="RHipPitch")
-    l_hip_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="LHipPitch")
-    r_knee: Mapped[float] = mapped_column(Float, default=0.0, name="RKnee")
-    l_knee: Mapped[float] = mapped_column(Float, default=0.0, name="LKnee")
-    r_ankle_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="RAnklePitch")
-    l_ankle_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="LAnklePitch")
-    r_ankle_roll: Mapped[float] = mapped_column(Float, default=0.0, name="RAnkleRoll")
-    l_ankle_roll: Mapped[float] = mapped_column(Float, default=0.0, name="LAnkleRoll")
-    head_pan: Mapped[float] = mapped_column(Float, default=0.0, name="HeadPan")
-    head_tilt: Mapped[float] = mapped_column(Float, default=0.0, name="HeadTilt")
+    r_shoulder_pitch: Mapped[float] = mapped_column(Float, name="RShoulderPitch")
+    l_shoulder_pitch: Mapped[float] = mapped_column(Float, name="LShoulderPitch")
+    r_shoulder_roll: Mapped[float] = mapped_column(Float, name="RShoulderRoll")
+    l_shoulder_roll: Mapped[float] = mapped_column(Float, name="LShoulderRoll")
+    r_elbow: Mapped[float] = mapped_column(Float, name="RElbow")
+    l_elbow: Mapped[float] = mapped_column(Float, name="LElbow")
+    r_hip_yaw: Mapped[float] = mapped_column(Float, name="RHipYaw")
+    l_hip_yaw: Mapped[float] = mapped_column(Float, name="LHipYaw")
+    r_hip_roll: Mapped[float] = mapped_column(Float, name="RHipRoll")
+    l_hip_roll: Mapped[float] = mapped_column(Float, name="LHipRoll")
+    r_hip_pitch: Mapped[float] = mapped_column(Float, name="RHipPitch")
+    l_hip_pitch: Mapped[float] = mapped_column(Float, name="LHipPitch")
+    r_knee: Mapped[float] = mapped_column(Float, name="RKnee")
+    l_knee: Mapped[float] = mapped_column(Float, name="LKnee")
+    r_ankle_pitch: Mapped[float] = mapped_column(Float, name="RAnklePitch")
+    l_ankle_pitch: Mapped[float] = mapped_column(Float, name="LAnklePitch")
+    r_ankle_roll: Mapped[float] = mapped_column(Float, name="RAnkleRoll")
+    l_ankle_roll: Mapped[float] = mapped_column(Float, name="LAnkleRoll")
+    head_pan: Mapped[float] = mapped_column(Float, name="HeadPan")
+    head_tilt: Mapped[float] = mapped_column(Float, name="HeadTilt")
 
     recording: Mapped["Recording"] = relationship("Recording", back_populates="joint_states")
 
@@ -203,26 +203,26 @@ class JointCommands(Base):
     _id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
     stamp: Mapped[float] = mapped_column(Float, nullable=False)
     recording_id: Mapped[int] = mapped_column(Integer, ForeignKey("Recording._id"), nullable=False)
-    r_shoulder_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="RShoulderPitch")
-    l_shoulder_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="LShoulderPitch")
-    r_shoulder_roll: Mapped[float] = mapped_column(Float, default=0.0, name="RShoulderRoll")
-    l_shoulder_roll: Mapped[float] = mapped_column(Float, default=0.0, name="LShoulderRoll")
-    r_elbow: Mapped[float] = mapped_column(Float, default=0.0, name="RElbow")
-    l_elbow: Mapped[float] = mapped_column(Float, default=0.0, name="LElbow")
-    r_hip_yaw: Mapped[float] = mapped_column(Float, default=0.0, name="RHipYaw")
-    l_hip_yaw: Mapped[float] = mapped_column(Float, default=0.0, name="LHipYaw")
-    r_hip_roll: Mapped[float] = mapped_column(Float, default=0.0, name="RHipRoll")
-    l_hip_roll: Mapped[float] = mapped_column(Float, default=0.0, name="LHipRoll")
-    r_hip_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="RHipPitch")
-    l_hip_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="LHipPitch")
-    r_knee: Mapped[float] = mapped_column(Float, default=0.0, name="RKnee")
-    l_knee: Mapped[float] = mapped_column(Float, default=0.0, name="LKnee")
-    r_ankle_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="RAnklePitch")
-    l_ankle_pitch: Mapped[float] = mapped_column(Float, default=0.0, name="LAnklePitch")
-    r_ankle_roll: Mapped[float] = mapped_column(Float, default=0.0, name="RAnkleRoll")
-    l_ankle_roll: Mapped[float] = mapped_column(Float, default=0.0, name="LAnkleRoll")
-    head_pan: Mapped[float] = mapped_column(Float, default=0.0, name="HeadPan")
-    head_tilt: Mapped[float] = mapped_column(Float, default=0.0, name="HeadTilt")
+    r_shoulder_pitch: Mapped[float] = mapped_column(Float, name="RShoulderPitch")
+    l_shoulder_pitch: Mapped[float] = mapped_column(Float, name="LShoulderPitch")
+    r_shoulder_roll: Mapped[float] = mapped_column(Float, name="RShoulderRoll")
+    l_shoulder_roll: Mapped[float] = mapped_column(Float, name="LShoulderRoll")
+    r_elbow: Mapped[float] = mapped_column(Float, name="RElbow")
+    l_elbow: Mapped[float] = mapped_column(Float, name="LElbow")
+    r_hip_yaw: Mapped[float] = mapped_column(Float, name="RHipYaw")
+    l_hip_yaw: Mapped[float] = mapped_column(Float, name="LHipYaw")
+    r_hip_roll: Mapped[float] = mapped_column(Float, name="RHipRoll")
+    l_hip_roll: Mapped[float] = mapped_column(Float, name="LHipRoll")
+    r_hip_pitch: Mapped[float] = mapped_column(Float, name="RHipPitch")
+    l_hip_pitch: Mapped[float] = mapped_column(Float, name="LHipPitch")
+    r_knee: Mapped[float] = mapped_column(Float, name="RKnee")
+    l_knee: Mapped[float] = mapped_column(Float, name="LKnee")
+    r_ankle_pitch: Mapped[float] = mapped_column(Float, name="RAnklePitch")
+    l_ankle_pitch: Mapped[float] = mapped_column(Float, name="LAnklePitch")
+    r_ankle_roll: Mapped[float] = mapped_column(Float, name="RAnkleRoll")
+    l_ankle_roll: Mapped[float] = mapped_column(Float, name="LAnkleRoll")
+    head_pan: Mapped[float] = mapped_column(Float, name="HeadPan")
+    head_tilt: Mapped[float] = mapped_column(Float, name="HeadTilt")
 
     recording: Mapped["Recording"] = relationship("Recording", back_populates="joint_commands")
 

ddlitlab2024/utils/utils.py

@@ -51,7 +51,7 @@ def shift_radian_to_positive_range(radian: float) -> float:
     :param radian: The pricipal range radian radian [-pi, pi].
     :return: The positive principal range radian [0, 2pi].
     """
-    return (radian + 2 * np.pi) % (2 * np.pi)
+    return (radian + 3 * np.pi) % (2 * np.pi)
 
 
 def timestamp_in_ns(seconds: int, nanoseconds: int) -> int: