google-deepmind/mujoco

[MJX] jax.lax.while_loop in solver.py prevents computation of backward gradients

Open

#2,259 opened on Nov 29, 2024

View on GitHub
 (14 comments) (5 reactions) (2 assignees)C++ (1,565 forks)github user discovery
MJXenhancementgood first issue

Repository metrics

Stars
 (13,782 stars)
PR merge metrics
 (PR metrics pending)

Description

The feature, motivation and pitch

Problem

The solver's jax.lax.while_loop implementation prevents gradient computation through the environment step during gradient based trajectory optimization. This occurs in the solver implementation when iterations > 1.

Error encountered with jax.jit compiled grad function:

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values.

Current workaround of using opt.iteration=1 leads to potentially inaccurate simulation and gradients.

Proposed Solution

Add an option to set a fixed iteration count (e.g., 4) that would be compatible with reverse-mode differentiation using either lax.scan or lax.fori_loop with static bounds.

Alternatives

No response

Additional context

No response

Contributor guide

Research direction
Replace the jax.lax.while loop in solver.py with a fixed iteration loop using lax.fori loop or lax.scan to enable reverse mode differentiation.
Tech stack
python
Domain
machine learning
Issue type
Bug
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Clear
Prerequisites
PythonJAX
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
60