# Copyright 2018 The dm_control Authors.

#

# 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 composer.Entity."""

import itertools

from absl.testing import absltest

from absl.testing import parameterized

from dm_control import mjcf

from dm_control.composer import arena

from dm_control.composer import define

from dm_control.composer import entity

from dm_control.composer.observation.observable import base as observable

import numpy as np

_NO_ROTATION = (1, 0, 0, 0) # Tests support for non-arrays and non-floats.

_NINETY_DEGREES_ABOUT_X = np.array(

[np.cos(np.pi / 4), np.sin(np.pi / 4), 0., 0.])

_NINETY_DEGREES_ABOUT_Y = np.array(

[np.cos(np.pi / 4), 0., np.sin(np.pi / 4), 0.])

_NINETY_DEGREES_ABOUT_Z = np.array(

[np.cos(np.pi / 4), 0., 0., np.sin(np.pi / 4)])

_FORTYFIVE_DEGREES_ABOUT_X = np.array(

[np.cos(np.pi / 8), np.sin(np.pi / 8), 0., 0.])

_TEST_ROTATIONS = [

# Triplets of original rotation, new rotation and final rotation.

(None, _NO_ROTATION, _NO_ROTATION),

(_NO_ROTATION, _NINETY_DEGREES_ABOUT_Z, _NINETY_DEGREES_ABOUT_Z),

(_FORTYFIVE_DEGREES_ABOUT_X, _NINETY_DEGREES_ABOUT_Y,

np.array([0.65328, 0.2706, 0.65328, -0.2706])),

]

def _param_product(**param_lists):

keys, values = zip(*param_lists.items())

for combination in itertools.product(*values):

yield dict(zip(keys, combination))

class TestEntity(entity.Entity):

"""Simple test entity that does nothing but declare some observables."""

def _build(self, name='test_entity'):

self._mjcf_root = mjcf.element.RootElement(model=name)

self._mjcf_root.worldbody.add('geom', type='sphere', size=(0.1,))

def _build_observables(self):

return TestEntityObservables(self)

@property

def mjcf_model(self):

return self._mjcf_root

class TestEntityObservables(entity.Observables):

"""Trivial observables for the test entity."""

@define.observable

def observable0(self):

return observable.Generic(lambda phys: 0.0)

@define.observable

def observable1(self):

return observable.Generic(lambda phys: 1.0)

class EntityTest(parameterized.TestCase):

def setUp(self):

super().setUp()

self.entity = TestEntity()

def testNumObservables(self):

"""Tests that the observables dict has the right number of entries."""

self.assertLen(self.entity.observables.as_dict(), 2)

def testObservableNames(self):

"""Tests that the observables dict keys correspond to the observable names.

"""

obs = self.entity.observables.as_dict()

self.assertIn('observable0', obs)

self.assertIn('observable1', obs)

subentity = TestEntity(name='subentity')

self.entity.attach(subentity)

self.assertIn('subentity/observable0', subentity.observables.as_dict())

self.assertEqual(subentity.observables.dict_keys.observable0,

'subentity/observable0')

self.assertIn('observable0', dir(subentity.observables.dict_keys))

self.assertIn('subentity/observable1', subentity.observables.as_dict())

self.assertEqual(subentity.observables.dict_keys.observable1,

'subentity/observable1')

self.assertIn('observable1', dir(subentity.observables.dict_keys))

def testEnableDisableObservables(self):

"""Test the enabling and disable functionality for observables."""

all_obs = self.entity.observables.as_dict()

self.entity.observables.enable_all()

for obs in all_obs.values():

self.assertTrue(obs.enabled)

self.entity.observables.disable_all()

for obs in all_obs.values():

self.assertFalse(obs.enabled)

self.entity.observables.observable0.enabled = True

self.assertTrue(all_obs['observable0'].enabled)

def testObservableDefaultOptions(self):

corruptor = lambda x: x

options = {

'update_interval': 2,

'buffer_size': 10,

'delay': 1,

'aggregator': 'max',

'corruptor': corruptor,

'enabled': True

}

self.entity.observables.set_options(options)

for obs in self.entity.observables.as_dict().values():

self.assertEqual(obs.update_interval, 2)

self.assertEqual(obs.delay, 1)

self.assertEqual(obs.buffer_size, 10)

self.assertEqual(obs.aggregator, observable.AGGREGATORS['max'])

self.assertEqual(obs.corruptor, corruptor)

self.assertTrue(obs.enabled)

def testObservablePartialDefaultOptions(self):

options = {'update_interval': 2, 'delay': 1}

self.entity.observables.set_options(options)

for obs in self.entity.observables.as_dict().values():

self.assertEqual(obs.update_interval, 2)

self.assertEqual(obs.delay, 1)

self.assertIsNone(obs.buffer_size)

self.assertIsNone(obs.aggregator)

self.assertIsNone(obs.corruptor)

def testObservableOptionsInvalidName(self):

options = {'asdf': None}

with self.assertRaisesRegex(KeyError, 'No observable with name \'asdf\''):

self.entity.observables.set_options(options)

def testObservableInvalidOptions(self):

options = {'observable0': {'asdf': 2}}

with self.assertRaisesRegex(AttributeError,

'Cannot add attribute asdf in configure.'):

self.entity.observables.set_options(options)

def testObservableOptions(self):

options = {

'observable0': {

'update_interval': 2,

'delay': 3

},

'observable1': {

'update_interval': 4,

'delay': 5

}

}

self.entity.observables.set_options(options)

observables = self.entity.observables.as_dict()

self.assertEqual(observables['observable0'].update_interval, 2)

self.assertEqual(observables['observable0'].delay, 3)

self.assertIsNone(observables['observable0'].buffer_size)

self.assertIsNone(observables['observable0'].aggregator)

self.assertIsNone(observables['observable0'].corruptor)

self.assertFalse(observables['observable0'].enabled)

self.assertEqual(observables['observable1'].update_interval, 4)

self.assertEqual(observables['observable1'].delay, 5)

self.assertIsNone(observables['observable1'].buffer_size)

self.assertIsNone(observables['observable1'].aggregator)

self.assertIsNone(observables['observable1'].corruptor)

self.assertFalse(observables['observable1'].enabled)

def testObservableOptionsEntityConstructor(self):

options = {

'observable0': {

'update_interval': 2,

'delay': 3

},

'observable1': {

'update_interval': 4,

'delay': 5

}

}

ent = TestEntity(observable_options=options)

observables = ent.observables.as_dict()

self.assertEqual(observables['observable0'].update_interval, 2)

self.assertEqual(observables['observable0'].delay, 3)

self.assertIsNone(observables['observable0'].buffer_size)

self.assertIsNone(observables['observable0'].aggregator)

self.assertIsNone(observables['observable0'].corruptor)

self.assertFalse(observables['observable0'].enabled)

self.assertEqual(observables['observable1'].update_interval, 4)

self.assertEqual(observables['observable1'].delay, 5)

self.assertIsNone(observables['observable1'].buffer_size)

self.assertIsNone(observables['observable1'].aggregator)

self.assertIsNone(observables['observable1'].corruptor)

self.assertFalse(observables['observable1'].enabled)

def testObservablePartialOptions(self):

options = {'observable0': {'update_interval': 2, 'delay': 3}}

self.entity.observables.set_options(options)

observables = self.entity.observables.as_dict()

self.assertEqual(observables['observable0'].update_interval, 2)

self.assertEqual(observables['observable0'].delay, 3)

self.assertIsNone(observables['observable0'].buffer_size)

self.assertIsNone(observables['observable0'].aggregator)

self.assertIsNone(observables['observable0'].corruptor)

self.assertFalse(observables['observable0'].enabled)

self.assertEqual(observables['observable1'].update_interval, 1)

self.assertIsNone(observables['observable1'].delay)

self.assertIsNone(observables['observable1'].buffer_size)

self.assertIsNone(observables['observable1'].aggregator)

self.assertIsNone(observables['observable1'].corruptor)

self.assertFalse(observables['observable1'].enabled)

def testAttach(self):

entities = [TestEntity() for _ in range(4)]

entities[0].attach(entities[1])

entities[1].attach(entities[2])

entities[0].attach(entities[3])

self.assertIsNone(entities[0].parent)

self.assertIs(entities[1].parent, entities[0])

self.assertIs(entities[2].parent, entities[1])

self.assertIs(entities[3].parent, entities[0])

self.assertIsNone(entities[0].mjcf_model.parent_model)

self.assertIs(entities[1].mjcf_model.parent_model, entities[0].mjcf_model)

self.assertIs(entities[2].mjcf_model.parent_model, entities[1].mjcf_model)

self.assertIs(entities[3].mjcf_model.parent_model, entities[0].mjcf_model)

self.assertEqual(list(entities[0].iter_entities()), entities)

def testDetach(self):

entities = [TestEntity() for _ in range(4)]

entities[0].attach(entities[1])

entities[1].attach(entities[2])

entities[0].attach(entities[3])

entities[1].detach()

with self.assertRaisesRegex(RuntimeError, 'not attached'):

entities[1].detach()

self.assertIsNone(entities[0].parent)

self.assertIsNone(entities[1].parent)

self.assertIs(entities[2].parent, entities[1])

self.assertIs(entities[3].parent, entities[0])

self.assertIsNone(entities[0].mjcf_model.parent_model)

self.assertIsNone(entities[1].mjcf_model.parent_model)

self.assertIs(entities[2].mjcf_model.parent_model, entities[1].mjcf_model)

self.assertIs(entities[3].mjcf_model.parent_model, entities[0].mjcf_model)

self.assertEqual(list(entities[0].iter_entities()),

[entities[0], entities[3]])

def testIterEntitiesExcludeSelf(self):

entities = [TestEntity() for _ in range(4)]

entities[0].attach(entities[1])

entities[1].attach(entities[2])

entities[0].attach(entities[3])

self.assertEqual(

list(entities[0].iter_entities(exclude_self=True)), entities[1:])

def testGlobalVectorToLocalFrame(self):

parent = TestEntity()

parent.mjcf_model.worldbody.add(

'site', xyaxes=[0, 1, 0, -1, 0, 0]).attach(self.entity.mjcf_model)

physics = mjcf.Physics.from_mjcf_model(parent.mjcf_model)

# 3D vectors

np.testing.assert_allclose(

self.entity.global_vector_to_local_frame(physics, [0, 1, 0]),

[1, 0, 0], atol=1e-10)

np.testing.assert_allclose(

self.entity.global_vector_to_local_frame(physics, [-1, 0, 0]),

[0, 1, 0], atol=1e-10)

np.testing.assert_allclose(

self.entity.global_vector_to_local_frame(physics, [0, 0, 1]),

[0, 0, 1], atol=1e-10)

# 2D vectors; z-component is ignored

np.testing.assert_allclose(

self.entity.global_vector_to_local_frame(physics, [0, 1]),

[1, 0], atol=1e-10)

np.testing.assert_allclose(

self.entity.global_vector_to_local_frame(physics, [-1, 0]),

[0, 1], atol=1e-10)

def testGlobalMatrixToLocalFrame(self):

parent = TestEntity()

parent.mjcf_model.worldbody.add(

'site', xyaxes=[0, 1, 0, -1, 0, 0]).attach(self.entity.mjcf_model)

physics = mjcf.Physics.from_mjcf_model(parent.mjcf_model)

rotation_atob = np.array([[0, 1, 0], [0, 0, -1], [-1, 0, 0]])

ego_rotation_atob = np.array([[0, 0, -1], [0, -1, 0], [-1, 0, 0]])

np.testing.assert_allclose(

self.entity.global_xmat_to_local_frame(physics, rotation_atob),

ego_rotation_atob, atol=1e-10)

flat_rotation_atob = np.reshape(rotation_atob, -1)

flat_rotation_ego_atob = np.reshape(ego_rotation_atob, -1)

np.testing.assert_allclose(

self.entity.global_xmat_to_local_frame(

physics, flat_rotation_atob),

flat_rotation_ego_atob, atol=1e-10)

@parameterized.parameters(*_param_product(

position=[None, [1., 0., -1.]],

quaternion=[None, _FORTYFIVE_DEGREES_ABOUT_X, _NINETY_DEGREES_ABOUT_Z],

freejoint=[False, True],

))

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_allclose(physics.bind(frame).xpos, ground_truth_pos)

np.testing.assert_allclose(physics.bind(frame).xquat, ground_truth_quat)

@parameterized.parameters(*_param_product(

original_position=[[-2, -1, -1.], [1., 0., -1.]],

position=[None, [1., 0., -1.]],

original_quaternion=_TEST_ROTATIONS[0],

quaternion=_TEST_ROTATIONS[1],

expected_quaternion=_TEST_ROTATIONS[2],

freejoint=[False, True],

))

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)

@parameterized.parameters(False, True)

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.)

if __name__ == '__main__':

absltest.main()