# Copyright 2017 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.

# ============================================================================

"""Mujoco functions to support named indexing.

The Mujoco name structure works as follows:

In mjxmacro.h, each "X" entry denotes a type (a), a field name (b) and a list

of dimension size metadata (c) which may contain both numbers and names, for

example

X(int, name_bodyadr, nbody, 1) // or

X(mjtNum, body_pos, nbody, 3)

a b c ----->

The second declaration states that the field `body_pos` has type `mjtNum` and

dimension sizes `(nbody, 3)`, i.e. the first axis is indexed by body number.

These and other named dimensions are sized based on the loaded model. This

information is parsed and stored in `mjbindings.sizes`.

In mjmodel.h, the struct mjModel contains an array of element name addresses

for each size name.

int* name_bodyadr; // body name pointers (nbody x 1)

By iterating over each of these element name address arrays, we first obtain a

mapping from size names to a list of element names.

{'nbody': ['cart', 'pole'], 'njnt': ['free', 'ball', 'hinge'], ...}

In addition to the element names that are derived from the mjModel struct at

runtime, we also assign hard-coded names to certain dimensions where there is an

established naming convention (e.g. 'x', 'y', 'z' for dimensions that correspond

to Cartesian positions).

For some dimensions, a single element name maps to multiple indices within the

underlying field. For example, a single joint name corresponds to a variable

number of indices within `qpos` that depends on the number of degrees of freedom

associated with that joint type. These are referred to as "ragged" dimensions.

In such cases we determine the size of each named element by examining the

address arrays (e.g. `jnt_qposadr`), and construct a mapping from size name to

element sizes:

{'nq': [7, 3, 1], 'nv': [6, 3, 1], ...}

Given these two dictionaries, we then create an `Axis` instance for each size

name. These objects have a `convert_key_item` method that handles the conversion

from indexing expressions containing element names to valid numpy indices.

Different implementations of `Axis` are used to handle "ragged" and "non-ragged"

dimensions.

{'nbody': RegularNamedAxis(names=['cart', 'pole']),

'nq': RaggedNamedAxis(names=['free', 'ball', 'hinge'], sizes=[7, 4, 1])}

We construct this dictionary once using `make_axis_indexers`.

Finally, for each field we construct a `FieldIndexer` class. A `FieldIndexer`

instance encapsulates a field together with a list of `Axis` instances (one per

dimension), and implements the named indexing logic by calling their respective

`convert_key_item` methods.

Summary of terminology:

* _size name_ or _size_ A dimension size name, e.g. `nbody` or `ngeom`.

* _element name_ or _name_ A named element in a Mujoco model, e.g. 'cart' or

'pole'.

* _element index_ or _index_ The index of an element name, for a specific size

name.

"""

import abc

import collections

import weakref

from dm_control.mujoco.wrapper import util

from dm_control.mujoco.wrapper.mjbindings import sizes

import numpy as np

# Mapping from {size_name: address_field_name} for ragged dimensions.

_RAGGED_ADDRS = {

'nq': 'jnt_qposadr',

'nv': 'jnt_dofadr',

'na': 'actuator_actadr',

'nsensordata': 'sensor_adr',

'nnumericdata': 'numeric_adr',

}

# Names of columns.

_COLUMN_NAMES = {

'xyz': ['x', 'y', 'z'],

'quat': ['qw', 'qx', 'qy', 'qz'],

'mat': ['xx', 'xy', 'xz',

'yx', 'yy', 'yz',

'zx', 'zy', 'zz'],

'rgba': ['r', 'g', 'b', 'a'],

}

# Mapping from keys of _COLUMN_NAMES to sets of field names whose columns are

# addressable using those names.

_COLUMN_ID_TO_FIELDS = {

'xyz': set([

'body_pos',

'body_ipos',

'body_inertia',

'jnt_pos',

'jnt_axis',

'geom_size',

'geom_pos',

'site_size',

'site_pos',

'cam_pos',

'cam_poscom0',

'cam_pos0',

'light_pos',

'light_dir',

'light_poscom0',

'light_pos0',

'light_dir0',

'mesh_vert',

'mesh_normal',

'mocap_pos',

'xpos',

'xipos',

'xanchor',

'xaxis',

'geom_xpos',

'site_xpos',

'cam_xpos',

'light_xpos',

'light_xdir',

'subtree_com',

'wrap_xpos',

'subtree_linvel',

'subtree_angmom',

]),

'quat': set([

'body_quat',

'body_iquat',

'geom_quat',

'site_quat',

'cam_quat',

'mocap_quat',

'xquat',

]),

'mat': set([

'cam_mat0',

'xmat',

'ximat',

'geom_xmat',

'site_xmat',

'cam_xmat',

]),

'rgba': set([

'geom_rgba',

'site_rgba',

'skin_rgba',

'mat_rgba',

'tendon_rgba',

])

}

def _get_size_name_to_element_names(model):

"""Returns a dict that maps size names to element names.

Args:

model: An instance of `mjbindings.mjModelWrapper`.

Returns:

A `dict` mapping from a size name (e.g. `'nbody'`) to a list of element

names.

"""

names = model.names

size_name_to_element_names = {}

for field_name in dir(model.ptr):

if not _is_name_pointer(field_name):

continue

# Get addresses of element names in `model.names` array, e.g.

# field name: `name_nbodyadr` and name_addresses: `[86, 92, 101]`, and skip

# when there are no elements for this type in the model.

name_addresses = getattr(model, field_name).ravel()

if not name_addresses.size:

continue

# Get the element names.

element_names = []

for start_index in name_addresses:

end_index = names.find(b'\0', start_index)

name = names[start_index:end_index]

element_names.append(str(name, 'utf-8'))

# String identifier for the size of the first dimension, e.g. 'nbody'.

size_name = _get_size_name(field_name)

size_name_to_element_names[size_name] = element_names

# Add custom element names for certain columns.

for size_name, element_names in _COLUMN_NAMES.items():

size_name_to_element_names[size_name] = element_names

# "Ragged" axes inherit their element names from other "non-ragged" axes.

# For example, the element names for "nv" axis come from "njnt".

for size_name, address_field_name in _RAGGED_ADDRS.items():

donor = 'n' + address_field_name.split('_')[0]

if donor == 'nactuator':

donor = 'nu'

if donor in size_name_to_element_names:

size_name_to_element_names[size_name] = size_name_to_element_names[donor]

# Mocap bodies are a special subset of bodies.

mocap_body_names = [None] * model.nmocap

for body_id, body_name in enumerate(size_name_to_element_names['nbody']):

body_mocapid = model.body_mocapid[body_id]

if body_mocapid != -1:

mocap_body_names[body_mocapid] = body_name

assert None not in mocap_body_names

size_name_to_element_names['nmocap'] = mocap_body_names

return size_name_to_element_names

def _get_size_name_to_element_sizes(model):

"""Returns a dict that maps size names to element sizes for ragged axes.

Args:

model: An instance of `mjbindings.mjModelWrapper`.

Returns:

A `dict` mapping from a size name (e.g. `'nv'`) to a numpy array of element

sizes. Size names corresponding to non-ragged axes are omitted.

"""

size_name_to_element_sizes = {}

for size_name, address_field_name in _RAGGED_ADDRS.items():

addresses = getattr(model, address_field_name).ravel()

if size_name == 'na':

element_sizes = np.where(addresses == -1, 0, 1)

else:

total_length = getattr(model, size_name)

element_sizes = np.diff(np.r_[addresses, total_length])

size_name_to_element_sizes[size_name] = element_sizes

return size_name_to_element_sizes

def make_axis_indexers(model):

"""Returns a dict that maps size names to `Axis` indexers.

Args:

model: An instance of `mjbindings.MjModelWrapper`.

Returns:

A `dict` mapping from a size name (e.g. `'nbody'`) to an `Axis` instance.

"""

size_name_to_element_names = _get_size_name_to_element_names(model)

size_name_to_element_sizes = _get_size_name_to_element_sizes(model)

# Unrecognized size names are treated as unnamed axes.

axis_indexers = collections.defaultdict(UnnamedAxis)

for size_name in size_name_to_element_names:

element_names = size_name_to_element_names[size_name]

if size_name in _RAGGED_ADDRS:

element_sizes = size_name_to_element_sizes[size_name]

singleton = (size_name == 'na')

indexer = RaggedNamedAxis(element_names, element_sizes,

singleton=singleton)

else:

indexer = RegularNamedAxis(element_names)

axis_indexers[size_name] = indexer

return axis_indexers

def _is_name_pointer(field_name):

"""Returns True for name pointer field names such as `name_bodyadr`."""

# Denotes name pointer fields in mjModel.

prefix, suffix = 'name_', 'adr'

return field_name.startswith(prefix) and field_name.endswith(suffix)

def _get_size_name(field_name, struct_name='mjmodel'):

# Look up size name in metadata.

try:

return sizes.array_sizes[struct_name][field_name][0]

except KeyError:

# Special handling required for name pointers in mjModel.

if _is_name_pointer(field_name):

return 'n' + field_name.split('_')[1][:-3]

else:

raise

def _validate_key_item(key_item):

if isinstance(key_item, (list, np.ndarray)):

for sub in key_item:

_validate_key_item(sub) # Recurse into nested arrays and lists.

elif key_item is Ellipsis:

raise IndexError('Ellipsis indexing not supported.')

elif key_item is None:

raise IndexError('None indexing not supported.')

elif key_item in (b'', u''):

raise IndexError('Empty strings are not allowed.')

class Axis(metaclass=abc.ABCMeta):

"""Handles the conversion of named indexing expressions into numpy indices."""

@abc.abstractmethod

def convert_key_item(self, key_item):

"""Converts a (possibly named) indexing expression to a numpy index."""

class UnnamedAxis(Axis):

"""An object representing an axis where the elements are not named."""

def convert_key_item(self, key_item):

"""Validate the indexing expression and return it unmodified."""

_validate_key_item(key_item)

return key_item

class RegularNamedAxis(Axis):

"""Represents an axis where each named element has a fixed size of 1."""

def __init__(self, names):

"""Initializes a new `RegularNamedAxis` instance.

Args:

names: A list or array of element names.

"""

self._names = names

self._names_to_offsets = {name: offset

for offset, name in enumerate(names) if name}

def convert_key_item(self, key_item):

"""Converts a named indexing expression to a numpy-friendly index."""

_validate_key_item(key_item)

if isinstance(key_item, str):

key_item = self._names_to_offsets[util.to_native_string(key_item)]

elif isinstance(key_item, (list, np.ndarray)):

# Cast lists to numpy arrays.

key_item = np.asarray(key_item)

original_shape = key_item.shape

# We assume that either all or none of the items in the array are strings

# representing names. If there is a mix, we will let NumPy throw an error

# when trying to index with the returned item.

if isinstance(key_item.flat[0], str):

key_item = np.array([self._names_to_offsets[util.to_native_string(k)]

for k in key_item.flat])

# Ensure the output shape is the same as that of the input.

key_item.shape = original_shape

return key_item

@property

def names(self):

"""Returns a list of element names."""

return self._names

class RaggedNamedAxis(Axis):

"""Represents an axis where the named elements may vary in size."""

def __init__(self, element_names, element_sizes, singleton=False):

"""Initializes a new `RaggedNamedAxis` instance.

Args:

element_names: A list or array containing the element names.

element_sizes: A list or array containing the size of each element.

singleton: Whether to reduce singleton slices to scalars.

"""

names_to_slices = {}

names_to_indices = {}

offset = 0

for name, size in zip(element_names, element_sizes):

# Don't add unnamed elements to the dicts.

if name:

if size == 1 and singleton:

names_to_slices[name] = offset

else:

names_to_slices[name] = slice(offset, offset + size)

names_to_indices[name] = range(offset, offset + size)

offset += size

self._names = element_names

self._sizes = element_sizes

self._names_to_slices = names_to_slices

self._names_to_indices = names_to_indices

def convert_key_item(self, key_item):

"""Converts a named indexing expression to a numpy-friendly index."""

_validate_key_item(key_item)

if isinstance(key_item, str):

key_item = self._names_to_slices[util.to_native_string(key_item)]

elif isinstance(key_item, (list, np.ndarray)):

# We assume that either all or none of the items in the sequence are

# strings representing names. If there is a mix, we will let NumPy throw

# an error when trying to index with the returned key.

if isinstance(key_item[0], str):

new_key = []

for k in key_item:

idx = self._names_to_indices[util.to_native_string(k)]

if isinstance(idx, int):

new_key.append(idx)

else:

new_key.extend(idx)

key_item = new_key

return key_item

@property

def names(self):

"""Returns a list of element names."""

return self._names

Axes = collections.namedtuple('Axes', ['row', 'col'])

Axes.__new__.__defaults__ = (None,) # Default value for optional 'col' field

class FieldIndexer:

"""An array-like object providing named access to a field in a MuJoCo struct.

FieldIndexers expose the same attributes and methods as an `np.ndarray`.

They may be indexed with strings or lists of strings corresponding to element

names. They also support standard numpy indexing expressions, with the

exception of indices containing `Ellipsis` or `None`.

"""

__slots__ = ('_field_name', '_field', '_axes')

def __init__(self,

parent_struct,

field_name,

axis_indexers):

"""Initializes a new `FieldIndexer`.

Args:

parent_struct: Wrapped ctypes structure, as generated by `mjbindings`.

field_name: String containing field name in `parent_struct`.

axis_indexers: A list of `Axis` instances, one per dimension.

"""

self._field_name = field_name

self._field = weakref.proxy(getattr(parent_struct, field_name))

self._axes = Axes(*axis_indexers)

def __dir__(self):

# Enables IPython tab completion

return sorted(set(dir(type(self)) + dir(self._field)))

def __getattr__(self, name):

return getattr(self._field, name)

def _convert_key(self, key):

"""Convert a (possibly named) indexing expression to a valid numpy index."""

return_tuple = isinstance(key, tuple)

if not return_tuple:

key = (key,)

if len(key) > self._field.ndim:

raise IndexError('Index tuple has {} elements, but array has only {} '

'dimensions.'.format(len(key), self._field.ndim))

new_key = tuple(axis.convert_key_item(key_item)

for axis, key_item in zip(self._axes, key))

if not return_tuple:

new_key = new_key[0]

return new_key

def __getitem__(self, key):

"""Converts the key to a numeric index and returns the indexed array.

Args:

key: Indexing expression.

Raises:

IndexError: If an indexing tuple has too many elements, or if it contains

`Ellipsis`, `None`, or an empty string.

Returns:

The indexed array.

"""

return self._field[self._convert_key(key)]

def __setitem__(self, key, value):

"""Converts the key and assigns to the indexed array.

Args:

key: Indexing expression.

value: Value to assign.

Raises:

IndexError: If an indexing tuple has too many elements, or if it contains

`Ellipsis`, `None`, or an empty string.

"""

self._field[self._convert_key(key)] = value

@property

def axes(self):

"""A namedtuple containing the row and column indexers for this field."""

return self._axes

def __repr__(self):

"""Returns a pretty string representation of the `FieldIndexer`."""

def get_name_arr_and_len(dim_idx):

"""Returns a string array of element names and the max name length."""

axis = self._axes[dim_idx]

size = self._field.shape[dim_idx]

try:

name_len = max(len(name) for name in axis.names)

name_arr = np.zeros(size, dtype='S{}'.format(name_len))

for name in axis.names:

if name:

# Use the `Axis` object to convert the name into a numpy index, then

# use this index to write into name_arr.

name_arr[axis.convert_key_item(name)] = name

except AttributeError:

name_arr = np.zeros(size, dtype='S0') # An array of zero-length strings

name_len = 0

return name_arr, name_len

row_name_arr, row_name_len = get_name_arr_and_len(0)

if self._field.ndim > 1:

col_name_arr, col_name_len = get_name_arr_and_len(1)

else:

col_name_arr, col_name_len = np.zeros(1, dtype='S0'), 0

idx_len = int(np.log10(max(self._field.shape[0], 1))) + 1

cls_template = '{class_name:}({field_name:}):'

col_template = '{padding:}{col_names:}'

row_template = '{idx:{idx_len:}} {row_name:>{row_name_len:}} {row_vals:}'

lines = []

# Write the class name and field name.

lines.append(cls_template.format(class_name=self.__class__.__name__,

field_name=self._field_name))

# Write a header line containing the column names (if there are any).

if col_name_len:

col_width = max(col_name_len, 9) + 1

extra_indent = 4

padding = ' ' * (idx_len + row_name_len + extra_indent)

col_names = ''.join(

'{name:<{width:}}'

.format(name=util.to_native_string(name), width=col_width)

for name in col_name_arr)

lines.append(col_template.format(padding=padding, col_names=col_names))

# Write the row names (if there are any) and the formatted array values.

if not self._field.shape[0]:

lines.append('(empty)')

else:

for idx, row in enumerate(self._field):

row_vals = np.array2string(

np.atleast_1d(row),

suppress_small=True,

formatter={'float_kind': '{: < 9.3g}'.format})

lines.append(row_template.format(

idx=idx,

idx_len=idx_len,

row_name=util.to_native_string(row_name_arr[idx]),

row_name_len=row_name_len,

row_vals=row_vals))

return '\n'.join(lines)

def struct_indexer(struct, struct_name, size_to_axis_indexer):

"""Returns an object with a `FieldIndexer` attribute for each dynamic field.

Usage example

```python

named_data = struct_indexer(mjdata, 'mjdata', size_to_axis_indexer)

fingertip_xpos = named_data.xpos['fingertip']

elbow_qvel = named_data.qvel['elbow']

```

Args:

struct: Wrapped ctypes structure as generated by `mjbindings`.

struct_name: String containing corresponding Mujoco name of struct.

size_to_axis_indexer: dict that maps size names to `Axis` instances.

Returns:

An object with a field for every dynamically sized array field, mapping to a

`FieldIndexer`. The returned object is immutable and has an `_asdict`

method.

Raises:

ValueError: If `struct_name` is not recognized.

"""

struct_name = struct_name.lower()

if struct_name not in sizes.array_sizes:

raise ValueError('Unrecognized struct name ' + struct_name)

array_sizes = sizes.array_sizes[struct_name]

# Used to create the namedtuple.

field_names = []

field_indexers = {}

for field_name in array_sizes:

# Skip over structured arrays and fields that have sizes but aren't numpy

# arrays, such as text fields and contacts (b/34805932).

attr = getattr(struct, field_name)

if not isinstance(attr, np.ndarray) or attr.dtype.fields:

continue

size_names = sizes.array_sizes[struct_name][field_name]

# Here we override the size name in order to enable named column indexing

# for certain fields, e.g. 3 becomes "xyz" for field name "xpos".

for new_col_size, field_set in _COLUMN_ID_TO_FIELDS.items():

if field_name in field_set:

size_names = (size_names[0], new_col_size)

break

axis_indexers = []

for size_name in size_names:

axis_indexers.append(size_to_axis_indexer[size_name])

field_indexers[field_name] = FieldIndexer(

parent_struct=struct,

field_name=field_name,

axis_indexers=axis_indexers)

field_names.append(field_name)

return make_struct_indexer(field_indexers)

def make_struct_indexer(field_indexers):

"""Returns an immutable container exposing named indexers as attributes."""

class StructIndexer:

__slots__ = ()

def _asdict(self):

return field_indexers.copy()

for name, indexer in field_indexers.items():

setattr(StructIndexer, name, indexer)

return StructIndexer()