Example of rolling out Python-based policies and MJX-based environments with passive viewer

[hartikainen]

Intro

Hello,

I'm a MuJoCo user working on manipulation.

My setup

mujoco==3.2.6
mujoco-mjx==3.2.6

My question

I'm looking for an example of how to sync mjpython passive viewer's state with external data. Specifically, I have an mjx-based environment and I want to visualize a rollout of trained jax-based policy while manually moving/perturbing the bodies from the passive viewer.

Say I have the following (somewhat incomplete and naive) passive viewer code:

policy = jax_based_policy()
environment = mjx_based_environment()
states = [environment_reset(reset_rng)]

with mujoco.viewer.launch_passive(model, data) as viewer:
  while viewer.is_running():
    action_rng, rng = jax.random.split(rng)
    action = policy_fn(action_rng, states[-1].obs)
    states.append(environment_step(states[-1], action))

    mjx.get_data_into(
        data,
        model,
        extract_first_state(states[-1].pipeline_state),
    )
    viewer.sync()
    time_until_next_step = model.opt.timestep - (time.time() - step_start)
    if time_until_next_step > 0:
        time.sleep(time_until_next_step)

This will obviously not work because none of the perturbations in data are synced into pipeline states. What should I sync from data into states[-1]? Do I need to do anything special here? Does there exist a clean example of handling such flow?

Minimal model and/or code that explain my question

No response

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.