Load URDF with frame IDs

crates/store/re_data_loader/src/loader_urdf.rs

@@ -12,7 +12,7 @@ use re_chunk::{ChunkBuilder, ChunkId, EntityPath, RowId, TimePoint};
 use re_log_types::{EntityPathPart, StoreId};
 use re_types::{
     AsComponents, Component as _, ComponentDescriptor, SerializedComponentBatch,
-    archetypes::{Asset3D, Transform3D},
+    archetypes::{Asset3D, CoordinateFrame, Transform3D},
     datatypes::Vec3D,
     external::glam,
 };
@@ -282,6 +282,15 @@ fn log_robot(
         .map(|prefix| prefix / EntityPath::from_single_string(urdf_tree.name.clone()))
         .unwrap_or_else(|| EntityPath::from_single_string(urdf_tree.name.clone()));
 
+    // The robot's root coordinate frame_id.
+    send_archetype(
+        tx,
+        store_id,
+        entity_path.clone(),
+        timepoint,
+        &CoordinateFrame::update_fields().with_frame(urdf_tree.root.name.clone()),
+    )?;
+
     walk_tree(
         &urdf_tree,
         tx,
@@ -345,8 +354,8 @@ fn log_joint(
         name: _,
         joint_type,
         origin,
-        parent: _,
-        child: _,
+        parent,
+        child,
         axis,
         limit,
         calibration,
@@ -355,7 +364,24 @@ fn log_joint(
         safety_controller,
     } = joint;
 
-    send_transform(tx, store_id, joint_path.clone(), origin, timepoint)?;
+    // A joint's own coordinate frame is that of its parent link.
+    send_archetype(
+        tx,
+        store_id,
+        joint_path.clone(),
+        timepoint,
+        &CoordinateFrame::update_fields().with_frame(parent.link.clone()),
+    )?;
+    // Send the joint origin, i.e. the default transform from parent link to child link.
+    send_transform(
+        tx,
+        store_id,
+        joint_path.clone(),
+        origin,
+        timepoint,
+        parent.link.clone(),
+        child.link.clone(),
+    )?;
 
     log_debug_format(
         tx,
@@ -404,13 +430,19 @@ fn log_joint(
     Ok(())
 }
 
-fn transform_from_pose(origin: &urdf_rs::Pose) -> Transform3D {
+fn transform_from_pose(
+    origin: &urdf_rs::Pose,
+    parent_frame: String,
+    child_frame: String,
+) -> Transform3D {
     let urdf_rs::Pose { xyz, rpy } = origin;
     let translation = [xyz[0] as f32, xyz[1] as f32, xyz[2] as f32];
     let quaternion = quat_xyzw_from_roll_pitch_yaw(rpy[0] as f32, rpy[1] as f32, rpy[2] as f32);
     Transform3D::update_fields()
         .with_translation(translation)
         .with_quaternion(quaternion)
+        .with_parent_frame(parent_frame)
+        .with_child_frame(child_frame)
 }
 
 fn send_transform(
@@ -419,21 +451,16 @@ fn send_transform(
     entity_path: EntityPath,
     origin: &urdf_rs::Pose,
     timepoint: &TimePoint,
+    parent_frame: String,
+    child_frame: String,
 ) -> anyhow::Result<()> {
-    let urdf_rs::Pose { xyz, rpy } = origin;
-    let is_identity = xyz.0 == [0.0, 0.0, 0.0] && rpy.0 == [0.0, 0.0, 0.0];
-
-    if is_identity {
-        Ok(()) // avoid noise
-    } else {
-        send_archetype(
-            tx,
-            store_id,
-            entity_path,
-            timepoint,
-            &transform_from_pose(origin),
-        )
-    }
+    send_archetype(
+        tx,
+        store_id,
+        entity_path,
+        timepoint,
+        &transform_from_pose(origin, parent_frame, child_frame),
+    )
 }
 
 /// Log the given value using its `Debug` formatting.
@@ -485,15 +512,24 @@ fn log_link(
         timepoint,
     )?;
 
+    // Log coordinate frame ID of the link.
+    send_archetype(
+        tx,
+        store_id,
+        link_entity.clone(),
+        timepoint,
+        &CoordinateFrame::update_fields().with_frame(link.name.clone()),
+    )?;
+
     for (i, visual) in visual.iter().enumerate() {
         let urdf_rs::Visual {
             name,
             origin,
             geometry,
             material,
         } = visual;
-        let name = name.clone().unwrap_or_else(|| format!("visual_{i}"));
-        let vis_entity = link_entity / EntityPathPart::new(name);
+        let visual_name = name.clone().unwrap_or_else(|| format!("visual_{i}"));
+        let visual_entity = link_entity / EntityPathPart::new(visual_name.clone());
 
         // Prefer inline defined material properties if present, otherwise fall back to global material.
         let material = material.as_ref().and_then(|mat| {
@@ -504,10 +540,33 @@ fn log_link(
             }
         });
 
-        send_transform(tx, store_id, vis_entity.clone(), origin, timepoint)?;
+        send_transform(
+            tx,
+            store_id,
+            visual_entity.clone(),
+            origin,
+            timepoint,
+            link.name.clone(),
+            visual_name,
+        )?;
+
+        let coordinate_frame = CoordinateFrame::update_fields().with_frame(link.name.clone());
+        send_archetype(
+            tx,
+            store_id,
+            visual_entity.clone(),
+            timepoint,
+            &coordinate_frame,
+        )?;
 
         log_geometry(
-            urdf_tree, tx, store_id, vis_entity, geometry, material, timepoint,
+            urdf_tree,
+            tx,
+            store_id,
+            visual_entity,
+            geometry,
+            material,
+            timepoint,
         )?;
     }
 
@@ -517,10 +576,27 @@ fn log_link(
             origin,
             geometry,
         } = collision;
-        let name = name.clone().unwrap_or_else(|| format!("collision_{i}"));
-        let collision_entity = link_entity / EntityPathPart::new(name);
+        let collision_name = name.clone().unwrap_or_else(|| format!("collision_{i}"));
+        let collision_entity = link_entity / EntityPathPart::new(collision_name.clone());
+
+        send_transform(
+            tx,
+            store_id,
+            collision_entity.clone(),
+            origin,
+            timepoint,
+            link.name.clone(),
+            collision_name,
+        )?;
 
-        send_transform(tx, store_id, collision_entity.clone(), origin, timepoint)?;
+        let coordinate_frame = CoordinateFrame::update_fields().with_frame(link.name.clone());
+        send_archetype(
+            tx,
+            store_id,
+            collision_entity.clone(),
+            timepoint,
+            &coordinate_frame,
+        )?;
 
         log_geometry(
             urdf_tree,

examples/rust/animated_urdf/src/main.rs

@@ -20,7 +20,7 @@ fn main() -> anyhow::Result<()> {
     use clap::Parser as _;
     let args = Args::parse();
 
-    let (rec, _serve_guard) = args.rerun.init("rerun_example_clock")?;
+    let (rec, _serve_guard) = args.rerun.init("rerun_example_animated_urdf")?;
     run(&rec, &args)
 }
 
@@ -47,17 +47,16 @@ fn run(rec: &rerun::RecordingStream, _args: &Args) -> anyhow::Result<()> {
                     joint.limit.lower..=joint.limit.upper,
                 );
 
-                // NOTE: each joint already has a fixed origin pose (logged with the URDF file),
-                // and Rerun won't allow us to override or add to that transform here.
-                // So instead we apply the dynamic rotation to the child link of the joint:
-                let child_link = urdf.get_joint_child(joint);
-                let link_path = urdf.get_link_path(child_link);
+                // Rerun loads the URDF transforms with child/parent frame relations.
+                // In order to move a joint, we just need to log a new transform between two of those frames.
                 rec.log(
-                    link_path,
+                    "/transforms",
                     &rerun::Transform3D::from_rotation(rerun::RotationAxisAngle::new(
                         fixed_axis,
                         dynamic_angle as f32,
-                    )),
+                    ))
+                    .with_parent_frame(joint.parent.link.clone())
+                    .with_child_frame(joint.child.link.clone()),
                 )?;
             }
         }