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

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

"""Utilities that are helpful for implementing initializers."""

import collections

def _get_root_model(mjcf_elements):

root_model = mjcf_elements[0].root.root_model

for element in mjcf_elements:

if element.root.root_model != root_model:

raise ValueError('entities do not all belong to the same root model')

return root_model

class JointStaticIsolator:

"""Helper class that isolates a collection of MuJoCo joints from others.

An instance of this class is a context manager that caches the positions and

velocities of all non-isolated joints *upon construction*, and resets them to

their original state when the context exits.

"""

def __init__(self, physics, joints):

"""Initializes the joint isolator.

Args:

physics: An instance of `mjcf.Physics`.

joints: An iterable of `mjcf.Element` representing joints that may be

modified inside the context managed by this isolator.

"""

if not isinstance(joints, collections.abc.Iterable):

joints = [joints]

root_model = _get_root_model(joints)

other_joints = [joint for joint in root_model.find_all('joint')

if joint not in joints]

if other_joints:

self._other_joints_mj = physics.bind(other_joints)

self._initial_qpos = self._other_joints_mj.qpos.copy()

self._initial_qvel = self._other_joints_mj.qvel.copy()

else:

self._other_joints_mj = None

def __enter__(self):

pass

def __exit__(self, exc_type, exc_value, traceback):

del exc_type, exc_value, traceback # unused

if self._other_joints_mj:

self._other_joints_mj.qpos = self._initial_qpos

self._other_joints_mj.qvel = self._initial_qvel