From f8217041b2720bf67d72291afb38a7ee7b137a95 Mon Sep 17 00:00:00 2001
From: diegoferigo <diego.ferigo@iit.it>
Date: Thu, 9 May 2024 17:26:12 +0200
Subject: [PATCH] Read frames from the original model description

---
 src/jaxsim/parsers/kinematic_graph.py |  3 ++-
 src/jaxsim/parsers/rod/parser.py      | 29 +++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/jaxsim/parsers/kinematic_graph.py b/src/jaxsim/parsers/kinematic_graph.py
index 3e7c3e2ac..42705f7b4 100644
--- a/src/jaxsim/parsers/kinematic_graph.py
+++ b/src/jaxsim/parsers/kinematic_graph.py
@@ -121,7 +121,8 @@ def __post_init__(self):
         # Also here, we assume the model is fixed-base, therefore the first frame will
         # have last_link_idx + 1. These frames are not part of the physics model.
         for index, frame in enumerate(self.frames):
-            frame.index = index + len(self.link_names())
+            with frame.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
+                frame.index = int(index + len(self.link_names()))
 
         # Number joints so that their index matches their child link index
         links_dict = {l.name: l for l in iter(self)}
diff --git a/src/jaxsim/parsers/rod/parser.py b/src/jaxsim/parsers/rod/parser.py
index c2e396d57..32ba44967 100644
--- a/src/jaxsim/parsers/rod/parser.py
+++ b/src/jaxsim/parsers/rod/parser.py
@@ -25,6 +25,7 @@ class SDFData(NamedTuple):
 
     link_descriptions: List[descriptions.LinkDescription]
     joint_descriptions: List[descriptions.JointDescription]
+    frame_descriptions: List[descriptions.LinkDescription]
     collision_shapes: List[descriptions.CollisionShape]
 
     sdf_model: rod.Model | None = None
@@ -113,6 +114,23 @@ def extract_model_data(
     # Create a dictionary to find easily links
     links_dict: Dict[str, descriptions.LinkDescription] = {l.name: l for l in links}
 
+    # ============
+    # Parse frames
+    # ============
+
+    # Parse the frames (unconnected)
+    frames = [
+        descriptions.LinkDescription(
+            name=f.name,
+            mass=jnp.array(0.0, dtype=float),
+            inertia=jnp.zeros(shape=(3, 3)),
+            parent=links_dict[f.attached_to],
+            pose=f.pose.transform() if f.pose is not None else jnp.eye(4),
+        )
+        for f in sdf_model.frames()
+        if f.attached_to in links_dict
+    ]
+
     # =========================
     # Process fixed-base models
     # =========================
@@ -309,6 +327,7 @@ def extract_model_data(
         model_name=sdf_model.name,
         link_descriptions=links,
         joint_descriptions=joints,
+        frame_descriptions=frames,
         collision_shapes=collisions,
         fixed_base=sdf_model.is_fixed_base(),
         base_link_name=sdf_model.get_canonical_link(),
@@ -356,6 +375,16 @@ def build_model_description(
         ],
     )
 
+    # Keep only the frames whose parent link is part of the parsed model.
+    frames_with_existing_parent_link = [
+        f for f in sdf_data.frame_descriptions if f.parent.name in model
+    ]
+
+    # Include the SDF frames originally stored in the SDF.
+    model = dataclasses.replace(
+        model, frames=frames_with_existing_parent_link + model.frames
+    )
+
     # Store the parsed SDF tree as extra info
     model = dataclasses.replace(model, extra_info={"sdf_model": sdf_data.sdf_model})