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

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

"""Helpers for modifying a walker to match mocap data."""

from dm_control import composer

from dm_control import mjcf

from dm_control.locomotion.mocap import mocap_pb2

from dm_control.locomotion.walkers import rescale

import numpy as np

class WalkerInfo:

"""Encapsulates routines that modify a walker to match mocap data."""

def __init__(self, proto):

"""Initializes this object.

Args:

proto: A `mocap_pb2.Walker` protocol buffer.

"""

self._proto = proto

def check_walker_is_compatible(self, walker):

"""Checks whether a given walker is compatible with this `WalkerInfo`."""

mocap_model = getattr(walker, 'mocap_walker_model', None)

if mocap_model is not None and mocap_model != self._proto.model:

model_type_name = list(mocap_pb2.Walker.Model.keys())[list(

mocap_pb2.Walker.Model.values()).index(self._proto.model)]

raise ValueError('Walker is not compatible with model type {!r}: got {}'

.format(model_type_name, walker))

def rescale_walker(self, walker):

"""Rescales a given walker to match the data in this `WalkerInfo`."""

self.check_walker_is_compatible(walker)

for subtree_info in self._proto.scaling.subtree:

body = walker.mjcf_model.find('body', subtree_info.body_name)

subtree_root = body.parent

if subtree_info.parent_length:

position_factor = subtree_info.parent_length / np.linalg.norm(body.pos)

else:

position_factor = subtree_info.size_factor

rescale.rescale_subtree(

subtree_root, position_factor, subtree_info.size_factor)

if self._proto.mass:

physics = mjcf.Physics.from_mjcf_model(walker.mjcf_model.root_model)

current_mass = physics.bind(walker.root_body).subtreemass

mass_factor = self._proto.mass / current_mass

for body in walker.root_body.find_all('body'):

inertial = getattr(body, 'inertial', None)

if inertial:

inertial.mass *= mass_factor

for geom in walker.root_body.find_all('geom'):

if geom.mass is not None:

geom.mass *= mass_factor

else:

current_density = geom.density if geom.density is not None else 1000

geom.density = current_density * mass_factor

def add_marker_sites(self, walker, size=0.01, rgba=(0., 0., 1., .3),

default_to_random_position=True, random_state=None):

"""Adds sites corresponding to mocap tracking markers."""

self.check_walker_is_compatible(walker)

random_state = random_state or np.random

sites = []

if self._proto.markers:

mocap_class = walker.mjcf_model.default.add('default', dclass='mocap')

mocap_class.site.set_attributes(type='sphere', size=(size,), rgba=rgba,

group=composer.SENSOR_SITES_GROUP)

for marker_info in self._proto.markers.marker:

body = walker.mjcf_model.find('body', marker_info.parent)

if not body:

raise ValueError('Walker model does not contain a body named {!r}'

.format(str(marker_info.parent)))

pos = marker_info.position

if not pos:

if default_to_random_position:

pos = random_state.uniform(-0.005, 0.005, size=3)

else:

pos = np.zeros(3)

sites.append(

body.add(

'site', name=str(marker_info.name), pos=pos, dclass=mocap_class))

walker.list_of_site_names = [site.name for site in sites]

return sites