def testShiftPose(self, original_position, position, original_quaternion,
quaternion, expected_quaternion, freejoint):
# Setup entity.
test_arena = arena.Arena()
subentity = TestEntity(name='subentity')
frame = test_arena.attach(subentity)
if freejoint:
frame.add('freejoint')
physics = mjcf.Physics.from_mjcf_model(test_arena.mjcf_model)
# Set the original position
subentity.set_pose(physics,
position=original_position,
quaternion=original_quaternion)
if position is None:
ground_truth_pos = original_position
else:
ground_truth_pos = original_position + np.array(position)
subentity.shift_pose(physics, position=position, quaternion=quaternion)
np.testing.assert_array_equal(
physics.bind(frame).xpos, ground_truth_pos)
updated_quat = physics.bind(frame).xquat
np.testing.assert_array_almost_equal(updated_quat, expected_quaternion,
1e-4)
def testShiftPoseWithVelocity(self, rotate_velocity):
# Setup entity.
test_arena = arena.Arena()
subentity = TestEntity(name='subentity')
frame = test_arena.attach(subentity)
frame.add('freejoint')
physics = mjcf.Physics.from_mjcf_model(test_arena.mjcf_model)
# Set the original position
subentity.set_pose(physics, position=[0., 0., 0.])
# Set velocity in y dim.
subentity.set_velocity(physics, [0., 1., 0.])
# Rotate the entity around the z axis.
subentity.shift_pose(physics,
quaternion=[0., 0., 0., 1.],
rotate_velocity=rotate_velocity)
physics.forward()
updated_position, _ = subentity.get_pose(physics)
if rotate_velocity:
# Should not have moved in the y dim.
np.testing.assert_array_almost_equal(updated_position[1], 0.)
else:
# Should not have moved in the x dim.
np.testing.assert_array_almost_equal(updated_position[0], 0.)
def testSetPose(self, position, quaternion, freejoint):
# Setup entity.
test_arena = arena.Arena()
subentity = TestEntity(name='subentity')
frame = test_arena.attach(subentity)
if freejoint:
frame.add('freejoint')
physics = mjcf.Physics.from_mjcf_model(test_arena.mjcf_model)
if quaternion is None:
ground_truth_quat = _NO_ROTATION
else:
ground_truth_quat = quaternion
if position is None:
ground_truth_pos = np.zeros(shape=(3, ))
else:
ground_truth_pos = position
subentity.set_pose(physics, position=position, quaternion=quaternion)
np.testing.assert_array_equal(
physics.bind(frame).xpos, ground_truth_pos)
np.testing.assert_array_equal(
physics.bind(frame).xquat, ground_truth_quat)
