Erik Frey
67fa7c1dc2
These io functions replace device_put, device_get_into. These new functions correctly copy intermediate fields such as efc_J, which the old functions could not. PiperOrigin-RevId: 586712273 Change-Id: Ia3a4104af17db2ebd41b41a2ca501c5de9ac5d05
229 lines
8.8 KiB
Python
229 lines
8.8 KiB
Python
# Copyright 2023 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Tests for scan functions."""
|
|
|
|
from absl.testing import absltest
|
|
from jax import numpy as jp
|
|
import mujoco
|
|
from mujoco import mjx
|
|
# pylint: disable=g-importing-member
|
|
from mujoco.mjx._src import scan
|
|
from mujoco.mjx._src.types import JointType
|
|
# pylint: enable=g-importing-member
|
|
import numpy as np
|
|
|
|
|
|
class ScanTest(absltest.TestCase):
|
|
|
|
_MULTI_DOF_XML = """
|
|
<mujoco>
|
|
<compiler inertiafromgeom="true"/>
|
|
<worldbody>
|
|
<body>
|
|
<joint type="free"/>
|
|
<geom size=".15" mass="1" type="sphere"/>
|
|
<body>
|
|
<joint axis="1 0 0" pos="1 0 0" type="ball"/>
|
|
<geom size=".15" mass="2" type="sphere"/>
|
|
</body>
|
|
<body>
|
|
<joint axis="1 0 0" pos="2 0 0" type="hinge"/>
|
|
<joint axis="0 1 0" pos="3 0 0" type="slide"/>
|
|
<geom size=".15" mass="3" type="sphere"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
def test_flat_empty(self):
|
|
"""Test scanning over just world body."""
|
|
m = mujoco.MjModel.from_xml_string("""
|
|
<mujoco model="world_body">
|
|
<worldbody/>
|
|
</mujoco>
|
|
""")
|
|
m = mjx.device_put(m)
|
|
|
|
def fn(body_id):
|
|
return body_id + 1
|
|
|
|
b_in = jp.array([1])
|
|
b_expect = jp.array([2])
|
|
b_out = scan.flat(m, fn, 'b', 'b', b_in)
|
|
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
def test_flat_joints(self):
|
|
"""Tests scanning over bodies with joints of different types."""
|
|
m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML)
|
|
m = mjx.device_put(m)
|
|
|
|
# we will test two functions:
|
|
# 1) j_fn receives jnt_types as a jp array
|
|
# 2) s_fn receives jnt_types as a static np array and can switch on it
|
|
j_fn = lambda jnt_pos, val: val + jp.sum(jnt_pos)
|
|
s_fn = lambda jnt_types, val: val + sum(jnt_types)
|
|
|
|
b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]])
|
|
b_expect = jp.array([[0, 0], [1, 1], [3, 3], [8, 8]])
|
|
b_out = scan.flat(m, j_fn, 'jb', 'b', m.jnt_pos, b_in)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
b_out = scan.flat(m, s_fn, 'jb', 'b', m.jnt_type, b_in)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
# None should be omitted from the results
|
|
def no_free(jnt_types, val):
|
|
if tuple(jnt_types) == (JointType.FREE,):
|
|
return None
|
|
return val + sum(jnt_types)
|
|
b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
|
|
b_out = scan.flat(m, no_free, 'jb', 'b', m.jnt_type, b_in)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
# we should not call functions for which we know we will discard the results
|
|
def no_world(jnt_types, val):
|
|
if jnt_types.size == 0:
|
|
self.fail('world has no dofs, should not be called')
|
|
return val + sum(jnt_types)
|
|
v_in = jp.ones((m.nv, 1))
|
|
scan.flat(m, no_world, 'jv', 'v', m.jnt_type, v_in)
|
|
|
|
def test_body_tree(self):
|
|
"""Tests tree scanning over bodies with different joint counts."""
|
|
m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML)
|
|
m = mjx.device_put(m)
|
|
|
|
# we will test two functions:
|
|
# 1) j_fn receives jnt_pos which is a jp array
|
|
# 2) s_fn receives jnt_types which is a static np array
|
|
def j_fn(carry, jnt_pos, val):
|
|
carry = jp.zeros_like(val) if carry is None else carry
|
|
return carry + val + jp.sum(jnt_pos)
|
|
|
|
def s_fn(carry, jnt_types, val):
|
|
carry = jp.zeros_like(val) if carry is None else carry
|
|
return carry + val + sum(jnt_types)
|
|
|
|
b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]])
|
|
b_expect = jp.array([[0, 0], [1, 1], [4, 4], [9, 9]])
|
|
|
|
b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
# and reverse too:
|
|
b_expect = jp.array([[12, 12], [12, 12], [3, 3], [8, 8]])
|
|
b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in, reverse=True)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in, reverse=True)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
# None should be omitted from the results
|
|
def no_free(carry, jnt_types, val):
|
|
if tuple(jnt_types) == (JointType.FREE,):
|
|
return None
|
|
carry = jp.zeros_like(val) if carry is None else carry
|
|
return carry + val + sum(jnt_types)
|
|
b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
|
|
b_out = scan.body_tree(m, no_free, 'jb', 'b', m.jnt_type, b_in)
|
|
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
|
|
|
|
_MULTI_ACT_XML = """
|
|
<mujoco>
|
|
<option timestep="0.02"/>
|
|
<compiler autolimits="true"/>
|
|
<default>
|
|
<geom contype="0" conaffinity="0"/>
|
|
</default>
|
|
<worldbody>
|
|
<body pos="0 0 -0.5">
|
|
<joint type="free" name="joint0" range="-37.81 86.15"/>
|
|
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
|
|
<body pos="0 0 -0.5">
|
|
<joint axis="1 0 0" type="hinge" name="joint1" range="-32.88 5.15" actuatorfrcrange="-0.48 0.72"/>
|
|
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
|
|
<body pos="0 0 -0.5">
|
|
<joint axis="0 1 0" type="slide" name="joint2" range="-83.31 54.77"/>
|
|
<joint axis="1 0 0" type="slide" name="joint3" range="-87.84 60.58" actuatorfrcrange="-0.62 0.28"/>
|
|
<joint axis="0 1 0" type="hinge" name="joint4" range="-49.96 83.03" actuatorfrcrange="-0.01 0.64"/>
|
|
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
|
|
<body pos="0 0 -0.5">
|
|
<joint axis="1 0 0" type="slide" name="joint5" range="-1.53 87.07" actuatorfrcrange="-0.06 0.26"/>
|
|
<joint axis="0 1 0" type="slide" name="joint6" range="-63.57 9.95"/>
|
|
<joint axis="1 0 0" type="hinge" name="joint7" range="-22.50 41.20" actuatorfrcrange="-0.41 0.84"/>
|
|
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
|
|
<body pos="0 0 -0.5">
|
|
<joint axis="0 1 0" type="hinge" name="joint8" range="-34.67 72.56"/>
|
|
<joint axis="1 0 0" type="hinge" name="joint9" range="-1.81 16.22" actuatorfrcrange="-0.41 0.14"/>
|
|
<joint axis="1 0 0" type="ball" name="joint10" range="0.00 6.02" actuatorfrcrange="-0.74 0.12"/>
|
|
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
<actuator>
|
|
<position joint="joint6" gear="1"/>
|
|
<intvelocity joint="joint1" kp="2000" actrange="-0.7 2.3"/>
|
|
<motor joint="joint3" gear="42"/>
|
|
<velocity joint="joint2" kv="123"/>
|
|
<position joint="joint0" ctrlrange="-0.9472 0.9472"/>
|
|
<intvelocity joint="joint7" kp="2000" actrange="-0.7 2.3"/>
|
|
<general joint="joint5" ctrlrange="0.1 2.34346" biastype="affine" gainprm="35 0 0" biasprm="0 -35 -0.65"/>
|
|
<general joint="joint4" ctrlrange="0.1 2.34346" biastype="affine" gainprm="35 0 0" biasprm="0 -35 -0.65"/>
|
|
</actuator>
|
|
</mujoco>
|
|
"""
|
|
|
|
def testscan_actuators(self):
|
|
"""Tests scanning over actuators."""
|
|
m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML)
|
|
m = mjx.device_put(m)
|
|
|
|
fn = lambda *args: args
|
|
args = (
|
|
m.actuator_gear,
|
|
m.jnt_type,
|
|
jp.arange(m.nq),
|
|
jp.arange(m.nv),
|
|
jp.array([1.4, 1.1]),
|
|
)
|
|
gear, jnt_typ, qadr, vadr, act = scan.flat(
|
|
m, fn, 'ujqva', 'ujqva', *args, group_by='u'
|
|
)
|
|
|
|
actuator_trnid = m.actuator_trnid[:, 0]
|
|
np.testing.assert_array_equal(gear, m.actuator_gear)
|
|
np.testing.assert_array_equal(jnt_typ, m.jnt_type[actuator_trnid])
|
|
np.testing.assert_array_equal(act, jp.array([1.4, 1.1]))
|
|
expected_vadr = np.concatenate(
|
|
[np.nonzero(m.dof_jntid == trnid)[0] for trnid in actuator_trnid]
|
|
)
|
|
np.testing.assert_array_equal(vadr, expected_vadr)
|
|
expected_qadr = np.concatenate(
|
|
[np.nonzero(scan._q_jointid(m) == i)[0] for i in actuator_trnid]
|
|
)
|
|
np.testing.assert_array_equal(qadr, expected_qadr)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
absltest.main()
|