Tendon/Muscle Support in Brax System

[ttktjmt]

Hello brax team,

I am exploring the integration of musculoskeletal model simulations using myosuite with brax, and I would like to see tendon and muscle features within the brax simulation system. Although brax now supports tendons and muscles on the backend via MJX compatibility, the visualization (for example, through the brax.io.html.render function) does not display tendons.

Issue

As I investigated the root cause of this problem, it seems that tendons are not currently implemented in the System class in brax/base.py and I cannot refer to them as objects from JavaScript files in brax/visualizer/js/*.js, which presumably contributes to this limitation. There may be other parts of the system that require modification as well.

I found this issue #149 really close to my request, but it is not active, probably because brax supported mjx and tendon was supported in the backend.

Request

Would it be possible to add this feature in development as a new feature for brax? It would significantly improve the visualization aspect of the simulation for researchers and developers working with musculoskeletal models.

If integrating this feature internally is not feasible at the moment, could you please advise on which files beyond brax/base.py and the visualization scripts in brax/visualizer/js/*.js would need modification for comprehensive tendon support?

Planned Implementation

• Visualize tendons as capsules, or with spheres at each node and a cylinder for the path (like mujocoUtils.js#L469-L483 and main.js#L221-L251 in mujoco_wasm)
• Display tendons in red color, with the color intensity/strength changing based on the activation strength of the attached muscle actuator.

Thank you in advance for considering this feature request.