Set some visual params as jax.Arrays for visual domain rando.

PiperOrigin-RevId: 692273879 Change-Id: I80475aa65744de06a6bfebbfb95a44faa751595f
google-deepmind · Nov 1, 2024 · 47ebb06 · 47ebb06
1 parent b6037d1
commit 47ebb06
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
@@ -963,8 +963,8 @@ class Model(PyTreeNode):
   light_bodyid: np.ndarray = _restricted_to('mujoco')
   light_targetbodyid: np.ndarray = _restricted_to('mujoco')
   light_directional: np.ndarray
-  light_pos: np.ndarray = _restricted_to('mujoco')
-  light_dir: np.ndarray = _restricted_to('mujoco')
+  light_pos: jax.Array
+  light_dir: jax.Array
   light_poscom0: np.ndarray = _restricted_to('mujoco')
   light_pos0: np.ndarray
   light_dir0: np.ndarray