naijekux
commented
6 months ago I suspect the NaN value error is raised from mjx.step() inside of brax.mjx.pipeline.step()
def step(
sys: System, state: State, act: jax.Array, unused_debug: bool = False
) -> State:
data = state.replace(ctrl=act)
data = mjx.step(sys, data)
q, qd = data.qpos, data.qvel
x = Transform(pos=data.xpos[1:], rot=data.xquat[1:])
cvel = Motion(vel=data.cvel[1:, 3:], ang=data.cvel[1:, :3])
offset = data.xpos[1:, :] - data.subtree_com[sys.body_rootid[1:]]
offset = Transform.create(pos=offset)
xd = offset.vmap().do(cvel)
data = _reformat_contact(sys, data)
return data.replace(q=q, qd=qd, x=x, xd=xd)for the reason that there's no problem by running brax.mjx.pipeline.init(), when the environment is being reset.
def init(
sys: System, q: jax.Array, qd: jax.Array, unused_debug: bool = False
) -> State:
data = mjx.make_data(sys)
data = data.replace(qpos=q, qvel=qd)
data = mjx.forward(sys, data)
q, qd = data.qpos, data.qvel
x = Transform(pos=data.xpos[1:], rot=data.xquat[1:])
cvel = Motion(vel=data.cvel[1:, 3:], ang=data.cvel[1:, :3])
offset = data.xpos[1:, :] - data.subtree_com[sys.body_rootid[1:]]
offset = Transform.create(pos=offset)
xd = offset.vmap().do(cvel)
data = _reformat_contact(sys, data)
return State(q=q, qd=qd, x=x, xd=xd, **data.__dict__)Apparently, the most codes in step() is the same as that in init() apart from the line data = mjx.step(sys, data). Now that the codes in init() generate no errors, NaN values come from calling mjx.step().
btaba
Hi,
I'm looking for some help with MJX.
The function
pipeline_step()('mjx' backend) inside ofstep()of gymnasium env framework doesn't work well and returns a new data (pipeline_state) consisting of NaN values.I found the issue by printing some variables inside of
step(), in order to test if every variables are valued well.and the output is like,
The only variable outputing NaN values at the beginning is
pipeline_state.q[:7], which is returned frompipeline_state = self.pipeline_step(state.pipeline_state, action), while the other variables were assigned before callingpipeline_step(). That means that the NaNs error happens inpipeline_step().I suppose that the two input parameters _pipelinestate and action have nothing with the problem mentioned above, but the MCJF model from a XML file causes the problem for it.
The same issue happens too when using brax 'generalized' instead of 'mjx' as backend. Details for it could be found here posted also in Brax's github
Here is the model created by a XML file:
Could someone please help me with this issue? Thanks