chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow jax config to choose precision.

Browse Source 
PiperOrigin-RevId: 606765036
Change-Id: I36977851615700507eb2e83cc6d2d1ff8255be62
This commit is contained in:
Baruch Tabanpour
2024-02-13 15:23:05 -08:00
committed by Copybara-Service
parent 9ecd700923
commit f50a44e193
2 changed files with 24 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+22 -22
View File
@@ -146,23 +146,23 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
ne, nf, nl, nc = constraint.count_constraints(m) ne, nf, nl, nc = constraint.count_constraints(m)
nefc = ne + nf + nl + nc nefc = ne + nf + nl + nc
zero_0 = jp.zeros(0, dtype=jp.float32) zero_0 = jp.zeros(0, dtype=float)
zero_nv = jp.zeros(m.nv, dtype=jp.float32) zero_nv = jp.zeros(m.nv, dtype=float)
zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32) zero_nv_6 = jp.zeros((m.nv, 6), dtype=float)
zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=jp.float32) zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=float)
zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32) zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=float)
zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32) zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=float)
zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32) zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=float)
zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32) zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=float)
zero_nefc = jp.zeros(nefc, dtype=jp.float32) zero_nefc = jp.zeros(nefc, dtype=float)
zero_na = jp.zeros(m.na, dtype=jp.float32) zero_na = jp.zeros(m.na, dtype=float)
zero_nu = jp.zeros(m.nu, dtype=jp.float32) zero_nu = jp.zeros(m.nu, dtype=float)
zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32) zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=float)
zero_nm = jp.zeros(m.nM, dtype=jp.float32) zero_nm = jp.zeros(m.nM, dtype=float)
# create first d to get num contacts and nc # create first d to get num contacts and nc
d = types.Data( d = types.Data(
solver_niter=jp.array(0, dtype=jp.int32), solver_niter=jp.array(0, dtype=int),
time=jp.array(0.0), time=jp.array(0.0),
qpos=jp.array(m.qpos0), qpos=jp.array(m.qpos0),
qvel=zero_nv, qvel=zero_nv,
@@ -171,31 +171,31 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
ctrl=zero_nu, ctrl=zero_nu,
qfrc_applied=zero_nv, qfrc_applied=zero_nv,
xfrc_applied=zero_nbody_6, xfrc_applied=zero_nbody_6,
eq_active=jp.zeros(m.neq, dtype=jp.int32), eq_active=jp.zeros(m.neq, dtype=int),
qacc=zero_nv, qacc=zero_nv,
act_dot=zero_na, act_dot=zero_na,
xpos=zero_nbody_3, xpos=zero_nbody_3,
xquat=jp.zeros((m.nbody, 4), dtype=jp.float32), xquat=jp.zeros((m.nbody, 4), dtype=float),
xmat=zero_nbody_3_3, xmat=zero_nbody_3_3,
xipos=zero_nbody_3, xipos=zero_nbody_3,
ximat=zero_nbody_3_3, ximat=zero_nbody_3_3,
xanchor=zero_njnt_3, xanchor=zero_njnt_3,
xaxis=zero_njnt_3, xaxis=zero_njnt_3,
geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32), geom_xpos=jp.zeros((m.ngeom, 3), dtype=float),
geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32), geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=float),
site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32), site_xpos=jp.zeros((m.nsite, 3), dtype=float),
site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32), site_xmat=jp.zeros((m.nsite, 3, 3), dtype=float),
subtree_com=zero_nbody_3, subtree_com=zero_nbody_3,
cdof=zero_nv_6, cdof=zero_nv_6,
cinert=zero_nbody_10, cinert=zero_nbody_10,
actuator_length=zero_nu, actuator_length=zero_nu,
actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), actuator_moment=jp.zeros((m.nu, m.nv), dtype=float),
crb=zero_nbody_10, crb=zero_nbody_10,
qM=zero_nm if support.is_sparse(m) else zero_nv_nv, qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLD=zero_nm if support.is_sparse(m) else zero_nv_nv, qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLDiagInv=zero_nv if support.is_sparse(m) else zero_0, qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
contact=types.Contact.zero(ncon), contact=types.Contact.zero(ncon),
efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32), efc_J=jp.zeros((nefc, m.nv), dtype=float),
efc_frictionloss=zero_nefc, efc_frictionloss=zero_nefc,
efc_D=zero_nefc, efc_D=zero_nefc,
actuator_velocity=zero_nu, actuator_velocity=zero_nu,
mjx/mujoco/mjx/_src/types.py
+2 -2
View File
@@ -579,8 +579,8 @@ class Contact(PyTreeNode):
solref=jp.zeros((ncon, mujoco.mjNREF)), solref=jp.zeros((ncon, mujoco.mjNREF)),
solreffriction=jp.zeros((ncon, mujoco.mjNREF)), solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
solimp=jp.zeros((ncon, mujoco.mjNIMP,)), solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
geom1=jp.zeros(ncon, dtype=jp.int32), geom1=jp.zeros(ncon, dtype=int),
geom2=jp.zeros(ncon, dtype=jp.int32), geom2=jp.zeros(ncon, dtype=int),
) )