From 68cd2633374fb6145648fae5e1fb64f15a5ae4f3 Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Mon, 20 Jan 2025 12:10:36 -0800
Subject: [PATCH] Change type of `light_poscom0` and `actuator_acc0'.

This is required for batching over `mjx.Model` in case recompilation is used to change inertial values.

PiperOrigin-RevId: 717596504
Change-Id: I4b6780903ebf12ebbf8b7bc7b8990a01609d4491
---
 mjx/mujoco/mjx/_src/types.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index 16fa3a81..030f4a27 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -999,7 +999,7 @@ class Model(PyTreeNode):
   light_castshadow: jax.Array
   light_pos: jax.Array
   light_dir: jax.Array
-  light_poscom0: np.ndarray = _restricted_to('mujoco')
+  light_poscom0: jax.Array
   light_pos0: np.ndarray
   light_dir0: np.ndarray
   light_cutoff: jax.Array
@@ -1129,7 +1129,7 @@ class Model(PyTreeNode):
   actuator_actrange: jax.Array
   actuator_gear: jax.Array
   actuator_cranklength: np.ndarray
-  actuator_acc0: np.ndarray
+  actuator_acc0: jax.Array
   actuator_lengthrange: np.ndarray
   actuator_plugin: np.ndarray = _restricted_to('mujoco')
   sensor_type: np.ndarray