Feature/ball joint support

[mohitgadde]

Summary

• Add expand_to_q_indices and expand_to_v_indices for ball joint indexing.
• Fix MDP observations, rewards, and events for ball joints.
• Works with Agility Digit-v3 robot configuration with ball joint.

Fixes #475

[NeoZng]

1. subtraction of quat is not defined, its result no long rep a rotation. (in joint_pos_rel() , class posture reward, and other joint reset events/reward/obs related impl)
2. in write_joint_position & write_joint_velocity, seems like they only support full quat & angular vel write.
3. default_root_state in Entity.py is not processed correctly. this gonna affect motion loader utility

[mohitgadde]

@NeoZng Thanks for the review.

Issue 1 (quaternion subtraction): Addressed. All quaternion difference computations now use quat_box_minus instead of direct subtraction:

• joint_pos_rel() in observations.py
• posture class in rewards.py
• variable_posture class in velocity/mdp/rewards.py
• reset_joints_by_offset in events.py (uses quat_from_angle_axis + quat_mul)

Issue 2 (write_joint_position/velocity): Could you clarify what's missing? I tested with Digit and selective writes work correctly - ball joints accept 4 values (quaternion) for position and 3 values (angular velocity) for velocity.

Issue 3 (default_root_state): Could you provide more context on what's incorrect? The shape differs by design (7 for fixed-base, 13 for floating-base), and events.py handles both cases. What specific scenario causes issues with the motion loader?

[NeoZng]

2: one must write all 4 qpos or 3 qvel for ball joint even just want to modify one. but this isnt a big deal.
3. my bad, it is a typo, i mean default_joint_pos and joint_pos . see mdp of tracking task. the robot has ball joints, the dim of the motion data [joint_pos] is inconsistent with joint_pos (most motion data only contain hinge joint pos)
4. new discover. i think ball joint would be a con for events.py:randomize_field().

maybe u should ask maintainer for more suggestions.

to support ball joints, the workaround I took was to treat the ball joints separately from the 1dof joints, just like what we do to floating joints. another ugly workaround is to use a series of hinges to replace ball joint, which works fine when the moving range is not that big (no gimbal lock)

[mohitgadde]

Thanks for the feedback!

Issue 2: Agreed.

Issues 3 & 4: I am aware of these edge cases. They are robot-specific scenarios:

• Issue 3 (motion data): Motion capture data typically only contains hinge joint positions. When the robot has ball joints, there's a dimension mismatch since ball joints need 4 qpos values (quaternion) vs 1 for hinge joints. This is inherently data/robot-specific.
• Issue 4 (randomization): Ball joints in closed-loop kinematic chains (parallel linkages) shouldn't have their initial orientations randomized - doing so would break the kinematic constraints and create invalid configurations. Ideally, a fixed pose with valid quaternions could be a way to go.

I would be happy to address Issues 3 & 4 as well if there is interest in supporting these scenarios more broadly.

Why this PR:

In my case to address ball-joints, I bypass Entity entirely and write directly to sim.data.qpos/sim.data.qvel. This means skipping Entity's joint state properties (joint_pos, joint_vel, default_joint_pos) and manually indexing with jnt_qposadr. Every MDP function that accesses joint data needs similar overrides - resets, observations, rewards, commands, etc. - and this becomes tedious.

This PR aims to fix that workaround by making Entity natively handle ball joints: correct indexing (4 qpos / 3 dof), proper q_offsets and v_offsets, and appropriate dimensions in joint_pos, joint_vel. Generic MDP functions would then work out of the box, with users only customizing truly robot-specific behavior (joint randomization, motion data formats, joint selection etc.).

I understand mjlab primarily supports robots with standard hinge joints (G1, Go2), so these edge cases may not be immediate priorities. That said, I would appreciate any thoughts on the best path forward.

@kevinzakka - looking forward to your feedback whenever you have a chance!

[mohitgadde]

@kevinzakka Hi! Just following up on this PR. I have resolved the merge conflicts with main and tests are passing. This addresses #475 (ball joint support). Would appreciate a review when you get a chance. Thanks!