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

Support for MJX bind set() for scalars.

Browse Source 
PiperOrigin-RevId: 715832281
Change-Id: I217021e81fae95e44e9297a1c447891ff76572bd
This commit is contained in:
Alessio Quaglino
2025-01-15 09:36:44 -08:00
committed by Copybara-Service
parent ff1ff44e34
commit f1d557c125
2 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support.py
+4
View File
@@ -437,6 +437,10 @@ class BindData(object):
def set(self, name: str, value: jax.Array) -> Data: def set(self, name: str, value: jax.Array) -> Data:
"""Set the value of an array in an MJX Data.""" """Set the value of an array in an MJX Data."""
array = getattr(self.data, self.__getname(name)) array = getattr(self.data, self.__getname(name))
try:
iter(value)
except TypeError:
value = [value]
if len(value) == 1: if len(value) == 1:
array = array.at[self.id].set(value[0]) array = array.at[self.id].set(value[0])
else: else:
mjx/mujoco/mjx/_src/support_test.py
+3
View File
@@ -236,6 +236,9 @@ class SupportTest(parameterized.TestCase):
dx4 = dx.bind(mx, s.actuators[1]).set('ctrl', [6]) dx4 = dx.bind(mx, s.actuators[1]).set('ctrl', [6])
np.testing.assert_array_equal(dx4.bind(mx, s.actuators).ctrl, [0, 6, 0]) np.testing.assert_array_equal(dx4.bind(mx, s.actuators).ctrl, [0, 6, 0])
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, [0, 0, 0]) np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, [0, 0, 0])
dx5 = dx.bind(mx, s.actuators[1]).set('ctrl', 7)
np.testing.assert_array_equal(dx5.bind(mx, s.actuators).ctrl, [0, 7, 0])
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, [0, 0, 0])
# test invalid name # test invalid name
with self.assertRaises(AttributeError): with self.assertRaises(AttributeError):