I am tring to run below code with MJX and jax. I wanna use GPU for running my trainig. when I run the official toturial code, I do not have any problem. I changed the model path with my model and then I have below Error.
The code :
import time
import itertools
import numpy as np
from typing import Callable, NamedTuple, Optional, Union, List
import mediapy as media
import matplotlib.pyplot as plt
# More legible printing from numpy.
np.set_printoptions(precision=3, suppress=True, linewidth=100)
from datetime import datetime
import functools
from IPython.display import HTML
import jax
from jax import numpy as jp
import numpy as np
from typing import Any, Dict, Sequence, Tuple, Union
from brax import base
from brax import envs
from brax import math
from brax.base import Base, Motion, Transform
from brax.envs.base import Env, PipelineEnv, State
from brax.mjx.base import State as MjxState
from brax.training.agents.ppo import train as ppo
from brax.training.agents.ppo import networks as ppo_networks
from brax.io import html, mjcf, model
from etils import epath
from flax import struct
from matplotlib import pyplot as plt
import mediapy as media
from ml_collections import config_dict
import mujoco
from mujoco import mjx
path = '/home/lab/Mohsen/gpu/UR5gripper_v3.xml'
with open (path, 'r') as file:
path = file.read()
from mujoco import mjx
# Make model, data, and renderer
mj_model = mujoco.MjModel.from_xml_string(path)
mj_data = mujoco.MjData(mj_model)
renderer = mujoco.Renderer(mj_model)
mjx_model = mjx.put_model(mj_model)
mjx_data = mjx.put_data(mj_model, mj_data)
print(mj_data.qpos, type(mj_data.qpos))
print(mjx_data.qpos, type(mjx_data.qpos), mjx_data.qpos.devices())
# enable joint visualization option:
scene_option = mujoco.MjvOption()
scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True
duration = 3.8 # (seconds)
framerate = 60 # (Hz)
frames = []
mujoco.mj_resetData(mj_model, mj_data)
while mj_data.time < duration:
mujoco.mj_step(mj_model, mj_data)
if len(frames) < mj_data.time * framerate:
renderer.update_scene(mj_data, scene_option=scene_option)
pixels = renderer.render()
frames.append(pixels)
# Simulate and display video.
media.write_video('/home/lab/Mohsen/gpu/first.mp4',frames, fps=framerate)
jit_step = jax.jit(mjx.step)
frames = []
mujoco.mj_resetData(mj_model, mj_data)
mjx_data = mjx.put_data(mj_model, mj_data)
while mjx_data.time < duration:
mjx_data = jit_step(mjx_model, mjx_data)
if len(frames) < mjx_data.time * framerate:
mj_data = mjx.get_data(mj_model, mjx_data)
renderer.update_scene(mj_data, scene_option=scene_option)
pixels = renderer.render()
frames.append(pixels)
media.write_video('/home/lab/Mohsen/gpu/gpu2.mp4',frames, fps=framerate)
the error is related to mjx_data = jit_step(mjx_model, mjx_data).
Error:
W external/tsl/tsl/framework/bfc_allocator.cc:291] Allocator (GPU_0_bfc) ran out of memory trying to allocate 18.38GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/lab/Mohsen/gpu/test.py", line 88, in <module>
mjx_data = jit_step(mjx_model, mjx_data)
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 19733914656 bytes.
then in terminal, I added export JAX_TRACEBACK_FILTERING=off
and now the error is changed to below:
Allocator (GPU_0_bfc) ran out of memory trying to allocate 18.38GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available.
Traceback (most recent call last):
File "/home/lab/Mohsen/gpu/test.py", line 88, in <module>
mjx_data = jit_step(mjx_model, mjx_data)
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 179, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/pjit.py", line 298, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr, attrs_tracked = _python_pjit_helper(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/pjit.py", line 176, in _python_pjit_helper
out_flat = pjit_p.bind(*args_flat, **params)
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/core.py", line 2788, in bind
return self.bind_with_trace(top_trace, args, params)
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/core.py", line 425, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/core.py", line 913, in process_primitive
return primitive.impl(*tracers, **params)
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/pjit.py", line 1488, in _pjit_call_impl
return xc._xla.pjit(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/pjit.py", line 1471, in call_impl_cache_miss
out_flat, compiled = _pjit_call_impl_python(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/pjit.py", line 1406, in _pjit_call_impl_python
lowering_parameters=mlir.LoweringParameters()).compile()
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2369, in compile
executable = UnloadedMeshExecutable.from_hlo(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2908, in from_hlo
xla_executable, compile_options = _cached_compilation(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2718, in _cached_compilation
xla_executable = compiler.compile_or_get_cached(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/compiler.py", line 333, in compile_or_get_cached
return _compile_and_write_cache(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/compiler.py", line 500, in _compile_and_write_cache
executable = backend_compile(
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/profiler.py", line 335, in wrapper
return func(*args, **kwargs)
File "/home/lab/.pyenv/versions/3.10.12/lib/python3.10/site-packages/jax/_src/compiler.py", line 238, in backend_compile
return backend.compile(built_c, compile_options=options)
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 19733914656 bytes.