google-deepmind/mujoco

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

Open

#2.259 aberto em 29 de nov. de 2024

Ver no GitHub
 (14 comments) (5 reactions) (2 assignees)C++ (1.565 forks)github user discovery
MJXenhancementgood first issue

Métricas do repositório

Stars
 (13.782 stars)
Métricas de merge de PR
 (Métricas PR pendentes)

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

Guia do colaborador

Direção de pesquisa
Substitua o jax.lax.while loop em solver.py por um loop de iteração fixa usando lax.fori loop ou lax.scan para habilitar a diferenciação de modo reverso.
Pilha de tecnologia
python
Domain
machine learning
Tipo Issue
Bug
DifficultyDificuldade de implementação estimada para um novo contribuidor, de 1 para alterações muito pequenas a 5 para trabalho de nível especializado.
3
Tempo estimadoUm intervalo de tempo aproximado para um colaborador experiente investigar, implementar, testar e preparar um pull request.
1-3 horas
Status da atividadeQuão disponível o issue aparece agora: novo, ativo, obsoleto, bloqueado ou aguardando entrada do mantenedor.
Ativo
ClarityCom que clareza o issue explica a mudança esperada, os critérios de aceitação e a próxima etapa.
Claro
Prerequisites
PythonJAX
Simpatia para novatosUma pontuação de 1 a 100 que estima o quão acessível este issue é para colaboradores iniciantes.
60