Add option to log F/T sensor data

Author: nepfaff

poetry.lock

@@ -7,11 +7,11 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.11"
-# drake = {version = ">=0.0.20240409", source = "drake-nightly"}
+# drake = {version = ">=0.0.20240302", source = "drake-nightly"}
 drake = "^1.28.0"
 numpy = "^1.25.2"
-manipulation = {git = "https://github.com/RussTedrake/manipulation.git", rev = "a2d393fdad5e6079db53b0c4e749e2a7fbc935b0"}
-iiwa_setup = {git = "https://github.com/nepfaff/iiwa_setup.git", rev="f0d5bbfc8378313bb88721e13cb909928e9ff31f"}
+manipulation = {git = "https://github.com/RussTedrake/manipulation.git", rev = "dd7851d2a8af76462a37f20604886872f56be5b0"}
+iiwa_setup = {git = "https://github.com/nepfaff/iiwa_setup.git", rev="5d1ae518e049443f97d088d9c8e787829ef073b9"}
 tqdm = "^4.66.1"
 hydra-core = "^1.3.2"
 omegaconf = "^2.3.0"

pyproject.toml

@@ -45,25 +45,33 @@ class JointData:
     ft_sensor_measurements: Optional[np.ndarray] = None
     """The force-torque sensor measurements of shape (T, 6) and form
     [f_x, f_y, f_z, tau_x, tau_y, tau_z]."""
+    ft_sensor_sample_times_s: Optional[np.ndarray] = None
+    """The sample times in seconds of the force-torque sensor measurements of shape
+    (T,)."""
 
     def remove_duplicate_samples(self) -> "JointData":
         """
         Removes duplicate samples from the joint data. Returns a new JointData object
         rather than modifying the current one.
         """
         _, unique_indices = np.unique(self.sample_times_s, return_index=True)
-        return JointData(
+        joint_data = JointData(
             joint_positions=self.joint_positions[unique_indices],
             joint_velocities=self.joint_velocities[unique_indices],
             joint_accelerations=self.joint_accelerations[unique_indices],
             joint_torques=self.joint_torques[unique_indices],
             sample_times_s=self.sample_times_s[unique_indices],
-            ft_sensor_measurements=(
-                None
-                if self.ft_sensor_measurements is None
-                else self.ft_sensor_measurements[unique_indices]
-            ),
         )
+        if self.ft_sensor_sample_times_s is not None:
+            _, unique_ft_indices = np.unique(
+                self.ft_sensor_sample_times_s, return_index=True
+            )
+            joint_data.ft_sensor_measurements = self.ft_sensor_measurements[
+                unique_ft_indices
+            ]
+            joint_data.ft_sensor_sample_times_s = self.ft_sensor_sample_times_s[
+                unique_ft_indices
+            ]
 
     def save_to_disk(self, path: Path) -> None:
         """Saves the joint data to disk.
@@ -79,6 +87,10 @@ def save_to_disk(self, path: Path) -> None:
         np.save(path / "sample_times_s.npy", self.sample_times_s)
         if self.ft_sensor_measurements is not None:
             np.save(path / "ft_sensor_measurements.npy", self.ft_sensor_measurements)
+        if self.ft_sensor_sample_times_s is not None:
+            np.save(
+                path / "ft_sensor_sample_times_s.npy", self.ft_sensor_sample_times_s
+            )
 
     @classmethod
     def load_from_disk(cls, path: Path) -> "JointData":
@@ -91,6 +103,7 @@ def load_from_disk(cls, path: Path) -> "JointData":
             The joint data.
         """
         ft_sensor_measurements_path = path / "ft_sensor_measurements.npy"
+        ft_sensor_sample_times_path = path / "ft_sensor_sample_times_s.npy"
         return cls(
             joint_positions=np.load(path / "joint_positions.npy"),
             joint_velocities=np.load(path / "joint_velocities.npy"),
@@ -102,6 +115,11 @@ def load_from_disk(cls, path: Path) -> "JointData":
                 if ft_sensor_measurements_path.exists()
                 else None
             ),
+            ft_sensor_sample_times_s=(
+                np.load(ft_sensor_sample_times_path)
+                if ft_sensor_sample_times_path.exists()
+                else None
+            ),
         )
 
     @classmethod

robot_payload_id/utils/dataclasses.py

@@ -6,10 +6,17 @@
 
 import numpy as np
 
+# export PYTHONPATH=~/src/ft_300s_driver/bazel-bin/lcmtypes/ft_300s/:${PYTHONPATH}
+try:
+    from ft_reading_t import ft_reading_t
+except ImportError:
+    logging.warning("Failed to import `ft_reading_t`. Won't be able to log FT data.")
+
 from iiwa_setup.controllers import (
     InverseDynamicsControllerWithGravityCompensationCancellation,
 )
 from iiwa_setup.iiwa import IiwaHardwareStationDiagram
+from iiwa_setup.sensors import FTSensorDataReceiver
 from manipulation.station import LoadScenario
 from pydrake.all import (
     ApplySimulatorConfig,
@@ -18,6 +25,9 @@
     ConstantVectorSource,
     Demultiplexer,
     DiagramBuilder,
+    DrakeLcm,
+    LcmInterfaceSystem,
+    LcmSubscriberSystem,
     MeshcatVisualizer,
     PiecewisePolynomial,
     Simulator,
@@ -107,6 +117,11 @@ def main():
         default=1e-3,
         help="The period at which to log data.",
     )
+    parser.add_argument(
+        "--log_ft_data",
+        action="store_true",
+        help="Whether to log the FT sensor data.",
+    )
     parser.add_argument(
         "--html_path",
         type=Path,
@@ -132,6 +147,7 @@ def main():
     duration_to_remove_at_start = args.duration_to_remove_at_start
     noise_scale = args.noise_scale
     logging_period = args.logging_period
+    log_ft_data = args.log_ft_data
     html_path = args.html_path
 
     assert not (use_hardware and noise_scale > 0.0)
@@ -251,6 +267,29 @@ def main():
             wsg_const_pos_source.get_output_port(), station.GetInputPort("wsg.position")
         )
 
+    if log_ft_data:
+        lcm = DrakeLcm()
+        lcm_system = builder.AddNamedSystem(
+            "lcm_interface_system", LcmInterfaceSystem(lcm)
+        )
+        ft_sensor_subscriber: LcmSubscriberSystem = builder.AddNamedSystem(
+            "ft_sensor_subscriber",
+            LcmSubscriberSystem.Make(
+                channel="FT300S_SENSOR_DATA",
+                lcm_type=ft_reading_t,
+                lcm=lcm_system,
+                use_cpp_serializer=False,
+                wait_for_message_on_initialization_timeout=10,
+            ),
+        )
+        ft_sensor_data_receiver: FTSensorDataReceiver = builder.AddNamedSystem(
+            "ft_sensor_data_receiver", FTSensorDataReceiver()
+        )
+        builder.Connect(
+            ft_sensor_subscriber.get_output_port(),
+            ft_sensor_data_receiver.get_input_port(),
+        )
+
     # Add data loggers
     desired_state_demux = builder.AddNamedSystem(
         "desired_state_demux", Demultiplexer(output_ports_sizes=[7, 7])
@@ -287,6 +326,12 @@ def main():
         "measured_torque_logger",
         VectorLogSink(num_positions, publish_period=logging_period),
     )
+    if log_ft_data:
+        # The FT 300-S sensor publishes at 100Hz.
+        ft_logger: VectorLogSink = builder.AddNamedSystem(
+            "ft_logger",
+            VectorLogSink(input_size=6, publish_period=np.min([1e-2, logging_period])),
+        )
     builder.Connect(
         desired_state_demux.get_output_port(0),
         commanded_position_logger.get_input_port(),
@@ -315,6 +360,11 @@ def main():
         station.GetOutputPort("iiwa.torque_measured"),
         measured_torque_logger.get_input_port(),
     )
+    if log_ft_data:
+        builder.Connect(
+            ft_sensor_data_receiver.GetOutputPort("ft_measured"),
+            ft_logger.get_input_port(),
+        )
 
     visualizer = MeshcatVisualizer.AddToBuilder(
         builder, station.GetOutputPort("query_object"), station.internal_meshcat
@@ -377,6 +427,10 @@ def main():
         simulator.get_context()
     ).sample_times()
 
+    if log_ft_data:
+        ft_data = ft_logger.FindLog(simulator.get_context()).data().T
+        ft_sample_times_s = ft_logger.FindLog(simulator.get_context()).sample_times()
+
     if only_log_excitation_traj_data:
         # Only keep data during excitation trajectory execution
         data_start_time = (
@@ -412,6 +466,14 @@ def main():
         # Shift sample times to start at 0
         sample_times_s -= sample_times_s[0]
 
+        if log_ft_data:
+            ft_start_idx = np.argmax(ft_sample_times_s >= data_start_time)
+            ft_end_idx = np.argmax(ft_sample_times_s >= excitation_traj_end_time)
+            ft_data = ft_data[ft_start_idx:ft_end_idx]
+            ft_sample_times_s = ft_sample_times_s[ft_start_idx:ft_end_idx]
+            # Shift sample times to start at 0
+            ft_sample_times_s -= ft_sample_times_s[0]
+
     # Remove duplicated samples
     _, unique_indices = np.unique(sample_times_s, return_index=True)
     commanded_position_data = commanded_position_data[unique_indices]
@@ -422,6 +484,10 @@ def main():
     commanded_torque_data = commanded_torque_data[unique_indices]
     measured_torque_data = measured_torque_data[unique_indices]
     sample_times_s = sample_times_s[unique_indices]
+    if log_ft_data:
+        _, ft_unique_indices = np.unique(ft_sample_times_s, return_index=True)
+        ft_data = ft_data[ft_unique_indices]
+        ft_sample_times_s = ft_sample_times_s[ft_unique_indices]
 
     # Add noise
     if noise_scale > 0.0:
@@ -446,6 +512,8 @@ def main():
             joint_accelerations=np.zeros_like(measured_position_data) * np.nan,
             joint_torques=measured_torque_data,
             sample_times_s=sample_times_s,
+            ft_sensor_measurements=ft_data if log_ft_data else None,
+            ft_sensor_sample_times_s=ft_sample_times_s if log_ft_data else None,
         )
         joint_data.save_to_disk(save_data_path)