Skip to content

Support USD parsing from physX fixed tendons and plumb through to MJWarp #654 #688

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
vreutskyy wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Support USD parsing from physX fixed tendons and plumb through to MJWarp #654 #688

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5327e7e
Support USD parsing from physX fixed tendons and plumb through to MJW…
vreutskyy Aug 29, 2025
0cbf77c
pre-commit fix
vreutskyy Aug 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 94 additions & 0 deletions newton/_src/sim/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -399,6 +399,13 @@ def __init__(self, up_axis: AxisType = Axis.Z, gravity: float = -9.81):
self.muscle_bodies = []
self.muscle_points = []

# tendons (fixed)
self.tendon_start = [] # start index in tendon_joints for each tendon
self.tendon_params = [] # (stiffness, damping, rest_length, lower_limit, upper_limit) for each tendon
self.tendon_joints = [] # joint indices attached to tendons
self.tendon_gearings = [] # gearing coefficients for each joint attachment
self.tendon_key = [] # string key for each tendon

# rigid bodies
self.body_mass = []
self.body_inertia = []
Expand Down Expand Up @@ -1174,11 +1181,31 @@ def transform_mul(a, b):
"equality_constraint_polycoef",
"equality_constraint_key",
"equality_constraint_enabled",
"tendon_params",
"tendon_key",
]

for attr in more_builder_attrs:
getattr(self, attr).extend(getattr(builder, attr))

# Handle tendon copying with proper joint indices offset
if builder.tendon_start:
# The joint offset to apply to tendon joint references
joint_offset = self.joint_count - builder.joint_count

# Update tendon_start indices to account for existing tendons
existing_tendon_joints = len(self.tendon_joints)
for i in range(len(builder.tendon_start)):
if i == 0:
self.tendon_start.append(existing_tendon_joints)
else:
self.tendon_start.append(existing_tendon_joints + builder.tendon_start[i])

# Copy tendon_joints with offset and tendon_gearings as-is
for joint_idx in builder.tendon_joints:
self.tendon_joints.append(joint_idx + joint_offset)
self.tendon_gearings.extend(builder.tendon_gearings)

self.joint_dof_count += builder.joint_dof_count
self.joint_coord_count += builder.joint_coord_count

Expand Down Expand Up @@ -2380,6 +2407,59 @@ def add_muscle(
# return the index of the muscle
return len(self.muscle_start) - 1

def add_tendon(
self,
name: str,
joint_ids: list[int],
gearings: list[float],
stiffness: float = 0.0,
damping: float = 0.0,
rest_length: float = 0.0,
lower_limit: float = float("-inf"),
upper_limit: float = float("inf"),
) -> int:
"""Adds a fixed tendon constraint between multiple joints.

Fixed tendons couple the motion of multiple joints through a linear
constraint on their positions. The tendon length is computed as:
L = rest_length + sum(gearing[i] * joint_pos[i])

Args:
name: A unique identifier for the tendon
joint_ids: List of joint indices that this tendon connects
gearings: Gearing coefficient for each joint (transmission ratio)
stiffness: Elastic stiffness of the tendon (0 for hard constraint)
damping: Damping coefficient
rest_length: Rest length of the tendon
lower_limit: Lower limit for tendon length
upper_limit: Upper limit for tendon length

Returns:
The index of the tendon in the model
"""
if len(joint_ids) != len(gearings):
raise ValueError("Number of joint IDs must match number of gearings")

if len(joint_ids) < 2:
raise ValueError("Tendon must connect at least 2 joints")

# Store the start index for this tendon's joint data
self.tendon_start.append(len(self.tendon_joints))

# Store tendon parameters
self.tendon_params.append((stiffness, damping, rest_length, lower_limit, upper_limit))

# Store the name/key
self.tendon_key.append(name)

# Store joint indices and gearings
for joint_id, gearing in zip(joint_ids, gearings, strict=False):
self.tendon_joints.append(joint_id)
self.tendon_gearings.append(gearing)

# Return the index of the tendon
return len(self.tendon_start) - 1

Comment on lines +2410 to +2462
Copy link
Contributor

@coderabbitai coderabbitai bot Aug 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

⚠️ Potential issue

add_tendon lacks essential validation (ids, dof, limits, degeneracy, duplicate name).

Invalid joint ids or non-scalar joints will explode later in the solver. Limits can be inverted; all-zero gearings yield a no-op tendon.

Apply:

 def add_tendon(
     self,
     name: str,
     joint_ids: list[int],
     gearings: list[float],
     stiffness: float = 0.0,
     damping: float = 0.0,
     rest_length: float = 0.0,
     lower_limit: float = float("-inf"),
     upper_limit: float = float("inf"),
 ) -> int:
@@
-        if len(joint_ids) != len(gearings):
+        if len(joint_ids) != len(gearings):
             raise ValueError("Number of joint IDs must match number of gearings")
 
-        if len(joint_ids) < 2:
+        if len(joint_ids) < 2:
             raise ValueError("Tendon must connect at least 2 joints")
 
+        # Basic parameter sanity
+        if lower_limit > upper_limit:
+            raise ValueError(f"lower_limit ({lower_limit}) must be <= upper_limit ({upper_limit})")
+        if name in self.tendon_key:
+            raise ValueError(f"Duplicate tendon name: {name!r}")
+        if len(set(joint_ids)) != len(joint_ids):
+            raise ValueError("Duplicate joint indices in tendon are not allowed")
+        if not any(abs(g) > 0.0 for g in gearings):
+            raise ValueError("All gearings are zero; tendon would be degenerate")
+
+        # Validate joint indices and that each is 1-DoF (revolute or prismatic)
+        invalid = [j for j in joint_ids if j < 0 or j >= self.joint_count]
+        if invalid:
+            raise IndexError(f"Invalid joint indices for tendon {name}: {invalid}")
+        for j in set(joint_ids):
+            lin, ang = self.joint_dof_dim[j]
+            if (lin + ang) != 1:
+                raise ValueError(
+                    f"Tendon {name} expects scalar joints; joint {j} has {lin+ang} DoFs"
+                )
+
         # Store the start index for this tendon's joint data
         self.tendon_start.append(len(self.tendon_joints))
@@
-        for joint_id, gearing in zip(joint_ids, gearings, strict=False):
+        for joint_id, gearing in zip(joint_ids, gearings, strict=False):
             self.tendon_joints.append(joint_id)
             self.tendon_gearings.append(gearing)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def add_tendon(
self,
name: str,
joint_ids: list[int],
gearings: list[float],
stiffness: float = 0.0,
damping: float = 0.0,
rest_length: float = 0.0,
lower_limit: float = float("-inf"),
upper_limit: float = float("inf"),
) -> int:
"""Adds a fixed tendon constraint between multiple joints.
Fixed tendons couple the motion of multiple joints through a linear
constraint on their positions. The tendon length is computed as:
L = rest_length + sum(gearing[i] * joint_pos[i])
Args:
name: A unique identifier for the tendon
joint_ids: List of joint indices that this tendon connects
gearings: Gearing coefficient for each joint (transmission ratio)
stiffness: Elastic stiffness of the tendon (0 for hard constraint)
damping: Damping coefficient
rest_length: Rest length of the tendon
lower_limit: Lower limit for tendon length
upper_limit: Upper limit for tendon length
Returns:
The index of the tendon in the model
"""
if len(joint_ids) != len(gearings):
raise ValueError("Number of joint IDs must match number of gearings")
if len(joint_ids) < 2:
raise ValueError("Tendon must connect at least 2 joints")
# Store the start index for this tendon's joint data
self.tendon_start.append(len(self.tendon_joints))
# Store tendon parameters
self.tendon_params.append((stiffness, damping, rest_length, lower_limit, upper_limit))
# Store the name/key
self.tendon_key.append(name)
# Store joint indices and gearings
for joint_id, gearing in zip(joint_ids, gearings, strict=False):
self.tendon_joints.append(joint_id)
self.tendon_gearings.append(gearing)
# Return the index of the tendon
return len(self.tendon_start) - 1
def add_tendon(
self,
name: str,
joint_ids: list[int],
gearings: list[float],
stiffness: float = 0.0,
damping: float = 0.0,
rest_length: float = 0.0,
lower_limit: float = float("-inf"),
upper_limit: float = float("inf"),
) -> int:
"""Adds a fixed tendon constraint between multiple joints.
Fixed tendons couple the motion of multiple joints through a linear
constraint on their positions. The tendon length is computed as:
L = rest_length + sum(gearing[i] * joint_pos[i])
Args:
name: A unique identifier for the tendon
joint_ids: List of joint indices that this tendon connects
gearings: Gearing coefficient for each joint (transmission ratio)
stiffness: Elastic stiffness of the tendon (0 for hard constraint)
damping: Damping coefficient
rest_length: Rest length of the tendon
lower_limit: Lower limit for tendon length
upper_limit: Upper limit for tendon length
Returns:
The index of the tendon in the model
"""
if len(joint_ids) != len(gearings):
raise ValueError("Number of joint IDs must match number of gearings")
if len(joint_ids) < 2:
raise ValueError("Tendon must connect at least 2 joints")
# Basic parameter sanity
if lower_limit > upper_limit:
raise ValueError(
f"lower_limit ({lower_limit}) must be <= upper_limit ({upper_limit})"
)
if name in self.tendon_key:
raise ValueError(f"Duplicate tendon name: {name!r}")
if len(set(joint_ids)) != len(joint_ids):
raise ValueError("Duplicate joint indices in tendon are not allowed")
if not any(abs(g) > 0.0 for g in gearings):
raise ValueError("All gearings are zero; tendon would be degenerate")
# Validate joint indices and that each is 1-DoF (revolute or prismatic)
invalid = [j for j in joint_ids if j < 0 or j >= self.joint_count]
if invalid:
raise IndexError(f"Invalid joint indices for tendon {name}: {invalid}")
for j in set(joint_ids):
lin, ang = self.joint_dof_dim[j]
if (lin + ang) != 1:
raise ValueError(
f"Tendon {name} expects scalar joints; joint {j} has {lin+ang} DoFs"
)
# Store the start index for this tendon's joint data
self.tendon_start.append(len(self.tendon_joints))
# Store tendon parameters
self.tendon_params.append(
(stiffness, damping, rest_length, lower_limit, upper_limit)
)
# Store the name/key
self.tendon_key.append(name)
# Store joint indices and gearings
for joint_id, gearing in zip(joint_ids, gearings, strict=False):
self.tendon_joints.append(joint_id)
self.tendon_gearings.append(gearing)
# Return the index of the tendon
return len(self.tendon_start) - 1

# region shapes

def add_shape(
Expand Down Expand Up @@ -4230,6 +4310,19 @@ def finalize(self, device: Devicelike | None = None, requires_grad: bool = False
m.muscle_points = wp.array(self.muscle_points, dtype=wp.vec3, requires_grad=requires_grad)
m.muscle_activations = wp.array(self.muscle_activations, dtype=wp.float32, requires_grad=requires_grad)

# -----------------------
# tendons

# close the tendon joint indices
tendon_start = copy.copy(self.tendon_start)
tendon_start.append(len(self.tendon_joints))

m.tendon_start = wp.array(tendon_start, dtype=wp.int32)
m.tendon_params = wp.array(self.tendon_params, dtype=wp.float32, requires_grad=requires_grad)
m.tendon_joints = wp.array(self.tendon_joints, dtype=wp.int32)
m.tendon_gearings = wp.array(self.tendon_gearings, dtype=wp.float32, requires_grad=requires_grad)
m.tendon_key = self.tendon_key

# --------------------------------------
# rigid bodies

Expand Down Expand Up @@ -4400,6 +4493,7 @@ def finalize(self, device: Devicelike | None = None, requires_grad: bool = False
m.edge_count = len(self.edge_rest_angle)
m.spring_count = len(self.spring_rest_length)
m.muscle_count = len(self.muscle_start)
m.tendon_count = len(self.tendon_start)
m.articulation_count = len(self.articulation_start)
m.equality_constraint_count = len(self.equality_constraint_type)

Expand Down
13 changes: 13 additions & 0 deletions newton/_src/sim/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,17 @@ def __init__(self, device: Devicelike | None = None):
self.muscle_activations = None
"""Muscle activations, shape [muscle_count], float."""

self.tendon_start = None
"""Start indices for each tendon's joint attachments, shape [tendon_count + 1], int."""
self.tendon_params = None
"""Tendon parameters (stiffness, damping, rest_length, lower_limit, upper_limit), shape [tendon_count, 5], float."""
self.tendon_joints = None
"""Joint indices that tendons attach to, int."""
self.tendon_gearings = None
"""Gearing coefficients for each joint attachment, float."""
self.tendon_key = []
"""Tendon keys/names, shape [tendon_count], str."""

self.body_q = None
"""Rigid body poses for state initialization, shape [body_count, 7], float."""
self.body_qd = None
Expand Down Expand Up @@ -357,6 +368,8 @@ def __init__(self, device: Devicelike | None = None):
"""Total number of springs in the system."""
self.muscle_count = 0
"""Total number of muscles in the system."""
self.tendon_count = 0
"""Total number of tendons in the system."""
self.articulation_count = 0
"""Total number of articulations in the system."""
self.joint_dof_count = 0
Expand Down
87 changes: 87 additions & 0 deletions newton/_src/solvers/mujoco/solver_mujoco.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2240,6 +2240,93 @@ def add_geoms(newton_body_id: int, incoming_xform: wp.transform | None = None):
eq.data[6:10] = wp.transform_get_rotation(eq_constraint_relpose[i])
eq.data[10] = eq_constraint_torquescale[i]

# Add fixed tendons from Newton model
if hasattr(model, "tendon_count") and model.tendon_count > 0:
# Access tendon data arrays on host
tendon_start = model.tendon_start.numpy() if hasattr(model.tendon_start, "numpy") else model.tendon_start
tendon_params = (
model.tendon_params.numpy() if hasattr(model.tendon_params, "numpy") else model.tendon_params
)
tendon_joints = (
model.tendon_joints.numpy() if hasattr(model.tendon_joints, "numpy") else model.tendon_joints
)
tendon_gearings = (
model.tendon_gearings.numpy() if hasattr(model.tendon_gearings, "numpy") else model.tendon_gearings
)

# Track tendon names to ensure uniqueness across environments
tendon_name_counts = {}

# Add each tendon
for i in range(model.tendon_count):
# Get the joints and gearings for this tendon
start_idx = tendon_start[i]
end_idx = tendon_start[i + 1]

# Check if all joints in this tendon are in the selected set
# Only add tendons where all joints are selected
all_joints_selected = True
for j in range(start_idx, end_idx):
joint_idx = tendon_joints[j]
if joint_idx not in selected_joints:
all_joints_selected = False
break

# Skip this tendon if not all joints are selected
if not all_joints_selected:
continue
Comment on lines +2268 to +2277
Copy link
Contributor

@coderabbitai coderabbitai bot Aug 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add bounds checking for joint indices.

The code should validate that joint indices are within valid bounds before using them to avoid potential crashes or undefined behavior.

 for j in range(start_idx, end_idx):
     joint_idx = tendon_joints[j]
+    if joint_idx < 0 or joint_idx >= model.joint_count:
+        print(f"Warning: Tendon {i} references invalid joint index {joint_idx}, skipping tendon")
+        all_joints_selected = False
+        break
     if joint_idx not in selected_joints:
         all_joints_selected = False
         break
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
all_joints_selected = True
for j in range(start_idx, end_idx):
joint_idx = tendon_joints[j]
if joint_idx not in selected_joints:
all_joints_selected = False
break
# Skip this tendon if not all joints are selected
if not all_joints_selected:
continue
all_joints_selected = True
for j in range(start_idx, end_idx):
joint_idx = tendon_joints[j]
if joint_idx < 0 or joint_idx >= model.joint_count:
print(f"Warning: Tendon {i} references invalid joint index {joint_idx}, skipping tendon")
all_joints_selected = False
break
if joint_idx not in selected_joints:
all_joints_selected = False
break
# Skip this tendon if not all joints are selected
if not all_joints_selected:
continue


# Create fixed tendon
tendon = spec.add_tendon()

# Ensure unique tendon name (similar to body naming)
base_name = model.tendon_key[i] if i < len(model.tendon_key) else f"tendon_{i}"
name = base_name
if name not in tendon_name_counts:
tendon_name_counts[name] = 1
else:
tendon_name_counts[name] += 1
name = f"{base_name}_{tendon_name_counts[name]}"

tendon.name = name

# Get tendon parameters (stiffness, damping, rest_length, lower_limit, upper_limit)
params = tendon_params[i]
stiffness = params[0]
damping = params[1]
rest_length = params[2]
lower_limit = params[3]
upper_limit = params[4]

# Set tendon properties
# In MuJoCo, fixed tendons use the following fields:
# - springlength: rest length of the tendon (needs to be array [2, 1])
# - stiffness: elastic stiffness
# - damping: damping coefficient
# - limited: whether limits are enabled
# - range: [lower_limit, upper_limit]
# springlength needs to be a 2x1 array for MuJoCo spec API
springlength_array = np.array([[rest_length], [0.0]], dtype=np.float64)
tendon.springlength = springlength_array
tendon.stiffness = stiffness
tendon.damping = damping

# Set limits if they are finite
if not (np.isinf(lower_limit) or np.isinf(upper_limit)):
tendon.limited = True
tendon.range[0] = lower_limit
tendon.range[1] = upper_limit

# Add joints to the fixed tendon
for j in range(start_idx, end_idx):
joint_idx = tendon_joints[j]
gearing = tendon_gearings[j]
joint_name = model.joint_key[joint_idx]

# In MuJoCo spec API, use wrap_joint to add a joint to the tendon
# This is the Python equivalent of mjs_wrapJoint in C API
tendon.wrap_joint(joint_name, gearing)
Comment on lines +2322 to +2328
Copy link
Contributor

@coderabbitai coderabbitai bot Aug 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add error handling for missing joint names.

The code assumes joint names exist in the model but doesn't handle the case where a joint might not have a name.

 for j in range(start_idx, end_idx):
     joint_idx = tendon_joints[j]
     gearing = tendon_gearings[j]
-    joint_name = model.joint_key[joint_idx]
+    joint_name = model.joint_key[joint_idx] if joint_idx < len(model.joint_key) else f"joint_{joint_idx}"
 
     # In MuJoCo spec API, use wrap_joint to add a joint to the tendon
     # This is the Python equivalent of mjs_wrapJoint in C API
-    tendon.wrap_joint(joint_name, gearing)
+    try:
+        tendon.wrap_joint(joint_name, gearing)
+    except Exception as e:
+        print(f"Warning: Failed to wrap joint '{joint_name}' to tendon '{name}': {e}")

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents 
In newton/_src/solvers/mujoco/solver_mujoco.py around lines 2322 to 2328, the
code dereferences model.joint_key[joint_idx] without checking that a name
exists; update it to validate that joint_idx is in range and that joint_name is
not None/empty before calling tendon.wrap_joint. If the name is missing, either
log a clear warning including joint_idx and tendon identifier and skip wrapping
that joint, or raise a descriptive exception if a missing name should be
considered fatal; ensure logging uses the existing logger and that control flow
continues safely when skipping.


assert len(spec.geoms) == colliding_shapes_per_env, (
"The number of geoms in the MuJoCo model does not match the number of colliding shapes in the Newton model."
)
Expand Down
Loading
Loading