Description
Env Setup:
NVIDIA A40 Driver Version: 550.90.07 CUDA Version: 12.5
python==3.11.12
jax==0.5.3
mujoco==3.3.1
brax==0.12.3
playground==0.0.4
When compiling and installing Madrona MJX and running the CartpoleBalance example with vision enabled, I encounter an internal cuSolver error originating in XLA. Curiously, initializing the environment once with vision disabled (so that Madrona’s physics backend compiles) and then re-initializing with vision enabled prevents the error.
Steps to Reproduce
- Install dependencies and compile Madrona MJX.
- Run the following snippet:
from mujoco_playground import dm_control_suite
from mujoco_playground import wrapper
import jax
from jax import numpy as jp
import time
num_envs = 4096
ctrl_dt = 0.04
episode_length = int(3 / ctrl_dt)
config_overrides = {
"vision": False,
"vision_config.render_batch_size": num_envs,
"action_repeat": 1,
"ctrl_dt": ctrl_dt,
"episode_length": episode_length,
}
env_name = "CartpoleBalance"
env = dm_control_suite.load(env_name, config_overrides=config_overrides)
config_overrides = {
"vision": True,
"vision_config.render_batch_size": num_envs,
"action_repeat": 1,
"ctrl_dt": ctrl_dt,
"episode_length": episode_length,
}
env_name = "CartpoleBalance"
env = dm_control_suite.load(env_name, config_overrides=config_overrides)
env = wrapper.wrap_for_brax_training(
env,
vision=True,
num_vision_envs=num_envs,
action_repeat=1,
episode_length=episode_length,
)
jit_reset = jax.jit(env.reset)
jit_step = jax.jit(env.step)
key_reset, key_act = jax.random.split(jax.random.PRNGKey(0))
state = jit_reset(jax.random.split(key_reset, num_envs))
N = 1000
img_shape = (64, 64)
t0 = time.time()
for i in range(N):
act = jax.random.uniform(
key_act, (num_envs, env.action_size), minval=-1.0, maxval=1.0
)
state = jit_step(state, act)
jax.tree_util.tree_map(
lambda x: x.block_until_ready(), state
) # Await device completion
dt = time.time() - t0
print("Madrona MJX: {:d} transitions per second".format(int(N * num_envs / dt)))Error message stack:
---------------------------------------------------------------------------
XlaRuntimeError Traceback (most recent call last)
Cell In[4], line 2
1 key_reset, key_act = jax.random.split(jax.random.PRNGKey(0))
----> 2 state = jit_reset(jax.random.split(key_reset, num_envs))
4 N = 1000
5 img_shape = (64, 64)
[... skipping hidden 11 frame]
File ~/miniforge3/envs/madmjx/lib/python3.11/site-packages/jax/_src/compiler.py:321, in backend_compile(backend, module, options, host_callbacks)
315 return backend.compile(
316 built_c, compile_options=options, host_callbacks=host_callbacks
317 )
318 # Some backends don't have `host_callbacks` option yet
319 # TODO(sharadmv): remove this fallback when all backends allow `compile`
320 # to take in `host_callbacks`
--> 321 return backend.compile(built_c, compile_options=options)
322 except xc.XlaRuntimeError as e:
323 for error_handler in _XLA_RUNTIME_ERROR_HANDLERS:
XlaRuntimeError: INTERNAL: cuSolver internal errorWhat I’ve Tried
- Isolating Brax vs. Madrona
- Setting "vision": False avoids the error entirely (physics backend compiles, steps run as expected).
- Two-stage initialization workaround
- Initialize and step through the env without vision.
- Re-initialize with vision.
This sequence successfully preloads the Madrona rendering backend and avoids the cuSolver error.
Expected Behavior
Environment should compile and run with vision enabled on the first initialization, without requiring a preliminary non-vision pass.
Possible Cause
Suspect that GPU VRAM preallocation or the initialization order of the physics vs. rendering backends leads to a cuSolver internal failure when vision is enabled on cold start.
Any insight into why the vision-enabled initialization triggers a cuSolver error on the first run?
Thank you!
Description
Env Setup:
When compiling and installing Madrona MJX and running the CartpoleBalance example with vision enabled, I encounter an internal cuSolver error originating in XLA. Curiously, initializing the environment once with vision disabled (so that Madrona’s physics backend compiles) and then re-initializing with vision enabled prevents the error.
Steps to Reproduce
Error message stack:
What I’ve Tried
This sequence successfully preloads the Madrona rendering backend and avoids the cuSolver error.
Expected Behavior
Environment should compile and run with vision enabled on the first initialization, without requiring a preliminary non-vision pass.
Possible Cause
Suspect that GPU VRAM preallocation or the initialization order of the physics vs. rendering backends leads to a cuSolver internal failure when vision is enabled on cold start.
Any insight into why the vision-enabled initialization triggers a cuSolver error on the first run?
Thank you!