From 47ebb0614e90ea99295b65ca383f04e564318b8c Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 1 Nov 2024 13:23:22 -0700 Subject: [PATCH] Set some visual params as jax.Arrays for visual domain rando. PiperOrigin-RevId: 692273879 Change-Id: I80475aa65744de06a6bfebbfb95a44faa751595f --- 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 e1df9681..ec17efed 100644 --- 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