Merge pull request #23 from thowell/io

update put_model

Author: thowell

mujoco/mjx/_src/forward.py

@@ -49,8 +49,8 @@ def next_activation(
 
     # get the high/low range for each actuator state
     limited = m.actuator_actlimited[actid]
-    range_low = wp.select(limited, -wp.inf, m.actuator_actrange[actid, 0])
-    range_high = wp.select(limited, wp.inf, m.actuator_actrange[actid, 1])
+    range_low = wp.select(limited, -wp.inf, m.actuator_actrange[actid][0])
+    range_high = wp.select(limited, wp.inf, m.actuator_actrange[actid][1])
 
     # get the actual actuation - skip if -1 (means stateless actuator)
     act_adr = m.actuator_actadr[actid]
@@ -65,7 +65,7 @@ def next_activation(
 
     # check dynType
     dyn_type = m.actuator_dyntype[actid]
-    dyn_prm = m.actuator_dynprm[actid, 0]
+    dyn_prm = m.actuator_dynprm[actid][0]
 
     # advance the actuation
     if dyn_type == 3:  # wp.static(WarpDynType.FILTEREXACT):

mujoco/mjx/_src/io.py

@@ -132,6 +132,11 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.dof_Madr = wp.array(mjm.dof_Madr, dtype=wp.int32, ndim=1)
   m.dof_armature = wp.array(mjm.dof_armature, dtype=wp.float32, ndim=1)
   m.dof_damping = wp.array(mjm.dof_damping, dtype=wp.float32, ndim=1)
+  m.actuator_actlimited = wp.array(mjm.actuator_actlimited, dtype=wp.int32, ndim=1)
+  m.actuator_actrange = wp.array(mjm.actuator_actrange, dtype=wp.vec2f, ndim=1)
+  m.actuator_actadr = wp.array(mjm.actuator_actadr, dtype=wp.int32, ndim=1)
+  m.actuator_dyntype = wp.array(mjm.actuator_dyntype, dtype=wp.int32, ndim=1)
+  m.actuator_dynprm = wp.array(mjm.actuator_dynprm, dtype=types.vec10f, ndim=1)
 
   return m
 

mujoco/mjx/_src/types.py

@@ -107,10 +107,10 @@ class Model:
   dof_armature: wp.array(dtype=wp.float32, ndim=1)
   dof_damping: wp.array(dtype=wp.float32, ndim=1)
   actuator_actlimited: wp.array(dtype=wp.int32, ndim=1)
-  actuator_actrange: wp.array(dtype=wp.float32, ndim=2)
+  actuator_actrange: wp.array(dtype=wp.vec2f, ndim=1)
   actuator_actadr: wp.array(dtype=wp.int32, ndim=1)
   actuator_dyntype: wp.array(dtype=wp.int32, ndim=1)
-  actuator_dynprm: wp.array(dtype=wp.float32, ndim=2)
+  actuator_dynprm: wp.array(dtype=vec10f, ndim=1)
 
 
 @wp.struct