Skip to content

Commit

Permalink
Read frames from the original model description
Browse files Browse the repository at this point in the history
  • Loading branch information
diegoferigo committed May 20, 2024
1 parent c14205a commit f821704
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/jaxsim/parsers/kinematic_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)}
Expand Down
29 changes: 29 additions & 0 deletions src/jaxsim/parsers/rod/parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
# =========================
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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})

Expand Down

0 comments on commit f821704

Please sign in to comment.