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

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

"""Tasks involving assembly and/or disassembly of bricks."""

import collections

from absl import logging

from dm_control import composer

from dm_control.composer import initializers

from dm_control.composer import variation

from dm_control.composer.observation import observable

from dm_control.composer.variation import distributions

from dm_control.entities import props

from dm_control.manipulation.shared import arenas

from dm_control.manipulation.shared import cameras

from dm_control.manipulation.shared import constants

from dm_control.manipulation.shared import observations

from dm_control.manipulation.shared import registry

from dm_control.manipulation.shared import robots

from dm_control.manipulation.shared import tags

from dm_control.manipulation.shared import workspaces

from dm_control.mujoco.wrapper import mjbindings

from dm_control.utils import rewards

import numpy as np

mjlib = mjbindings.mjlib

_BrickWorkspace = collections.namedtuple(

'_BrickWorkspace',

['prop_bbox', 'tcp_bbox', 'goal_hint_pos', 'goal_hint_quat', 'arm_offset'])

# Ensures that the prop does not collide with the table during initialization.

_PROP_Z_OFFSET = 1e-6

_WORKSPACE = _BrickWorkspace(

prop_bbox=workspaces.BoundingBox(

lower=(-0.1, -0.1, _PROP_Z_OFFSET),

upper=(0.1, 0.1, _PROP_Z_OFFSET)),

tcp_bbox=workspaces.BoundingBox(

lower=(-0.1, -0.1, 0.15),

upper=(0.1, 0.1, 0.4)),

goal_hint_pos=(0.2, 0.1, 0.),

goal_hint_quat=(-0.38268343, 0., 0., 0.92387953),

arm_offset=robots.ARM_OFFSET)

# Alpha value of the visual goal hint representing the goal state for each task.

_HINT_ALPHA = 0.75

# Distance thresholds for the shaping rewards for getting the top brick close

# to the bottom brick, and for 'clicking' them together.

_CLOSE_THRESHOLD = 0.01

_CLICK_THRESHOLD = 0.001

# Sequence of colors for the brick(s).

_COLOR_VALUES, _COLOR_NAMES = list(

zip(

((1., 0., 0.), 'red'),

((0., 1., 0.), 'green'),

((0., 0., 1.), 'blue'),

((0., 1., 1.), 'cyan'),

((1., 0., 1.), 'magenta'),

((1., 1., 0.), 'yellow'),

))

class _Common(composer.Task):

"""Common components of brick tasks."""

def __init__(self,

arena,

arm,

hand,

num_bricks,

obs_settings,

workspace,

control_timestep):

if not 2 <= num_bricks <= 6:

raise ValueError('`num_bricks` must be between 2 and 6, got {}.'

.format(num_bricks))

if num_bricks > 3:

# The default values computed by MuJoCo's compiler are too small if there

# are more than three stacked bricks, since each stacked pair generates

# a large number of contacts. The values below are sufficient for up to

# 6 stacked bricks.

# TODO(b/78331644): It may be useful to log the size of `physics.model`

# and `physics.data` after compilation to gauge the

# impact of these changes on MuJoCo's memory footprint.

arena.mjcf_model.size.nconmax = 400

arena.mjcf_model.size.njmax = 1200

self._arena = arena

self._arm = arm

self._hand = hand

self._arm.attach(self._hand)

self._arena.attach_offset(self._arm, offset=workspace.arm_offset)

self.control_timestep = control_timestep

# Add custom camera observable.

self._task_observables = cameras.add_camera_observables(

arena, obs_settings, cameras.FRONT_CLOSE)

color_sequence = iter(_COLOR_VALUES)

brick_obs_options = observations.make_options(

obs_settings, observations.FREEPROP_OBSERVABLES)

bricks = []

brick_frames = []

goal_hint_bricks = []

for _ in range(num_bricks):

color = next(color_sequence)

brick = props.Duplo(color=color,

observable_options=brick_obs_options)

brick_frames.append(arena.add_free_entity(brick))

bricks.append(brick)

# Translucent, contactless brick with no observables. These are used to

# provide a visual hint representing the goal state for each task.

hint_brick = props.Duplo(color=color)

_hintify(hint_brick, alpha=_HINT_ALPHA)

arena.attach(hint_brick)

goal_hint_bricks.append(hint_brick)

self._bricks = bricks

self._brick_frames = brick_frames

self._goal_hint_bricks = goal_hint_bricks

# Position and quaternion for the goal hint.

self._goal_hint_pos = workspace.goal_hint_pos

self._goal_hint_quat = workspace.goal_hint_quat

self._tcp_initializer = initializers.ToolCenterPointInitializer(

self._hand, self._arm,

position=distributions.Uniform(*workspace.tcp_bbox),

quaternion=workspaces.DOWN_QUATERNION)

# Add sites for visual debugging.

workspaces.add_bbox_site(

body=self.root_entity.mjcf_model.worldbody,

lower=workspace.tcp_bbox.lower,

upper=workspace.tcp_bbox.upper,

rgba=constants.GREEN, name='tcp_spawn_area')

workspaces.add_bbox_site(

body=self.root_entity.mjcf_model.worldbody,

lower=workspace.prop_bbox.lower,

upper=workspace.prop_bbox.upper,

rgba=constants.BLUE, name='prop_spawn_area')

@property

def task_observables(self):

return self._task_observables

@property

def root_entity(self):

return self._arena

@property

def arm(self):

return self._arm

@property

def hand(self):

return self._hand

class Stack(_Common):

"""Build a stack of Duplo bricks."""

def __init__(self,

arena,

arm,

hand,

num_bricks,

target_height,

moveable_base,

randomize_order,

obs_settings,

workspace,

control_timestep):

"""Initializes a new `Stack` task.

Args:

arena: `composer.Entity` instance.

arm: `robot_base.RobotArm` instance.

hand: `robot_base.RobotHand` instance.

num_bricks: The total number of bricks; must be between 2 and 6.

target_height: The target number of bricks in the stack in order to get

maximum reward. Must be between 2 and `num_bricks`.

moveable_base: Boolean specifying whether or not the bottom brick should

be moveable.

randomize_order: Boolean specifying whether to randomize the desired order

of bricks in the stack at the start of each episode.

obs_settings: `observations.ObservationSettings` instance.

workspace: A `_BrickWorkspace` instance.

control_timestep: Float specifying the control timestep in seconds.

Raises:

ValueError: If `num_bricks` is not between 2 and 6, or if

`target_height` is not between 2 and `num_bricks - 1`.

"""

if not 2 <= target_height <= num_bricks:

raise ValueError('`target_height` must be between 2 and {}, got {}.'

.format(num_bricks, target_height))

super().__init__(

arena=arena,

arm=arm,

hand=hand,

num_bricks=num_bricks,

obs_settings=obs_settings,

workspace=workspace,

control_timestep=control_timestep)

self._moveable_base = moveable_base

self._randomize_order = randomize_order

self._target_height = target_height

self._prop_bbox = workspace.prop_bbox

# Shuffled at the start of each episode if `randomize_order` is True.

self._desired_order = np.arange(target_height)

# In the random order case, create a `prop_pose` observable that informs the

# agent of the desired order.

if randomize_order:

desired_order_observable = observable.Generic(self._get_desired_order)

desired_order_observable.configure(**obs_settings.prop_pose._asdict())

self._task_observables['desired_order'] = desired_order_observable

def _get_desired_order(self, physics):

del physics # Unused

return self._desired_order.astype(np.double)

def initialize_episode_mjcf(self, random_state):

if self._randomize_order:

self._desired_order = random_state.choice(

len(self._bricks), size=self._target_height, replace=False)

logging.info('Desired stack order (from bottom to top): [%s]',

' '.join(_COLOR_NAMES[i] for i in self._desired_order))

# If the base of the stack should be fixed, remove the freejoint for the

# first brick (and ensure that all the others have freejoints).

fixed_indices = [] if self._moveable_base else [self._desired_order[0]]

_add_or_remove_freejoints(attachment_frames=self._brick_frames,

fixed_indices=fixed_indices)

# We need to define the prop initializer for the bricks here rather than in

# the `__init__`, since `PropPlacer` looks for freejoints on instantiation.

self._brick_placer = initializers.PropPlacer(

props=self._bricks,

position=distributions.Uniform(*self._prop_bbox),

quaternion=workspaces.uniform_z_rotation,

settle_physics=True)

def initialize_episode(self, physics, random_state):

self._brick_placer(physics, random_state)

self._hand.set_grasp(physics, close_factors=random_state.uniform())

self._tcp_initializer(physics, random_state)

# Arrange the goal hint bricks in the desired stack order.

_build_stack(physics,

bricks=self._goal_hint_bricks,

base_pos=self._goal_hint_pos,

base_quat=self._goal_hint_quat,

order=self._desired_order,

random_state=random_state)

def get_reward(self, physics):

pairs = list(zip(self._desired_order[:-1], self._desired_order[1:]))

pairwise_rewards = _get_pairwise_stacking_rewards(

physics=physics, bricks=self._bricks, pairs=pairs)

# The final reward is an average over the pairwise rewards.

return np.mean(pairwise_rewards)

class Reassemble(_Common):

"""Disassemble a stack of Duplo bricks and reassemble it in another order."""

def __init__(self,

arena,

arm,

hand,

num_bricks,

randomize_initial_order,

randomize_desired_order,

obs_settings,

workspace,

control_timestep):

"""Initializes a new `Reassemble` task.

Args:

arena: `composer.Entity` instance.

arm: `robot_base.RobotArm` instance.

hand: `robot_base.RobotHand` instance.

num_bricks: The total number of bricks; must be between 2 and 6.

randomize_initial_order: Boolean specifying whether to randomize the

initial order of bricks in the stack at the start of each episode.

randomize_desired_order: Boolean specifying whether to independently

randomize the desired order of bricks in the stack at the start of each

episode. By default the desired order will be the reverse of the initial

order, with the exception of the base brick which is always the same as

in the initial order since it is welded in place.

obs_settings: `observations.ObservationSettings` instance.

workspace: A `_BrickWorkspace` instance.

control_timestep: Float specifying the control timestep in seconds.

Raises:

ValueError: If `num_bricks` is not between 2 and 6.

"""

super().__init__(

arena=arena,

arm=arm,

hand=hand,

num_bricks=num_bricks,

obs_settings=obs_settings,

workspace=workspace,

control_timestep=control_timestep)

self._randomize_initial_order = randomize_initial_order

self._randomize_desired_order = randomize_desired_order

# Randomized at the start of each episode if `randomize_initial_order` is

# True.

self._initial_order = np.arange(num_bricks)

# Randomized at the start of each episode if `randomize_desired_order` is

# True.

self._desired_order = self._initial_order.copy()

self._desired_order[1:] = self._desired_order[-1:0:-1]

# In the random order case, create a `prop_pose` observable that informs the

# agent of the desired order.

if randomize_desired_order:

desired_order_observable = observable.Generic(self._get_desired_order)

desired_order_observable.configure(**obs_settings.prop_pose._asdict())

self._task_observables['desired_order'] = desired_order_observable

# Distributions of positions and orientations for the base of the stack.

self._base_pos = distributions.Uniform(*workspace.prop_bbox)

self._base_quat = workspaces.uniform_z_rotation

def _get_desired_order(self, physics):

del physics # Unused

return self._desired_order.astype(np.double)

def initialize_episode_mjcf(self, random_state):

if self._randomize_initial_order:

random_state.shuffle(self._initial_order)

# The bottom brick will be fixed to the table, so it must be the same in

# both the initial and desired order.

self._desired_order[0] = self._initial_order[0]

# By default the desired order of the other bricks is the opposite of their

# initial order.

self._desired_order[1:] = self._initial_order[-1:0:-1]

if self._randomize_desired_order:

random_state.shuffle(self._desired_order[1:])

logging.info('Desired stack order (from bottom to top): [%s]',

' '.join(_COLOR_NAMES[i] for i in self._desired_order))

# Remove the freejoint from the bottom brick in the stack.

_add_or_remove_freejoints(attachment_frames=self._brick_frames,

fixed_indices=[self._initial_order[0]])

def initialize_episode(self, physics, random_state):

# Build the initial stack.

_build_stack(physics,

bricks=self._bricks,

base_pos=self._base_pos,

base_quat=self._base_quat,

order=self._initial_order,

random_state=random_state)

# Arrange the goal hint bricks into a stack with the desired order.

_build_stack(physics,

bricks=self._goal_hint_bricks,

base_pos=self._goal_hint_pos,

base_quat=self._goal_hint_quat,

order=self._desired_order,

random_state=random_state)

self._hand.set_grasp(physics, close_factors=random_state.uniform())

self._tcp_initializer(physics, random_state)

def get_reward(self, physics):

pairs = list(zip(self._desired_order[:-1], self._desired_order[1:]))

# We set `close_coef=0.` because the coarse shaping reward causes problems

# for this task (it means there is a strong disincentive to break up the

# initial stack).

pairwise_rewards = _get_pairwise_stacking_rewards(

physics=physics,

bricks=self._bricks,

pairs=pairs,

close_coef=0.)

# The final reward is an average over the pairwise rewards.

return np.mean(pairwise_rewards)

def _distance(pos1, pos2):

diff = pos1 - pos2

return sum(np.sqrt((diff * diff).sum(1)))

def _min_stud_to_hole_distance(physics, bottom_brick, top_brick):

# Positions of the top left and bottom right studs on the `bottom_brick` and

# the top left and bottom right holes on the `top_brick`.

stud_pos = physics.bind(bottom_brick.studs[[0, -1], [0, -1]]).xpos

hole_pos = physics.bind(top_brick.holes[[0, -1], [0, -1]]).xpos

# Bricks are rotationally symmetric, so we compute top left -> top left and

# top left -> bottom right distances and return whichever of these is smaller.

dist1 = _distance(stud_pos, hole_pos)

dist2 = _distance(stud_pos[::-1], hole_pos)

return min(dist1, dist2)

def _get_pairwise_stacking_rewards(physics, bricks, pairs, close_coef=0.1):

"""Returns a vector of shaping reward components based on pairwise distances.

Args:

physics: An `mjcf.Physics` instance.

bricks: A list of `composer.Entity` instances corresponding to bricks.

pairs: A list of `(bottom_idx, top_idx)` tuples specifying which pairs of

bricks should be measured.

close_coef: Float specfying the relative weight given to the coarse-

tolerance shaping component for getting the bricks close to one another

(as opposed to the fine-tolerance component for clicking them together).

Returns:

A numpy array of size `len(pairs)` containing values in (0, 1], where

1 corresponds to a stacked pair of bricks.

"""

distances = []

for bottom_idx, top_idx in pairs:

bottom_brick = bricks[bottom_idx]

top_brick = bricks[top_idx]

distances.append(

_min_stud_to_hole_distance(physics, bottom_brick, top_brick))

distances = np.hstack(distances)

# Coarse-tolerance component for bringing the holes close to the studs.

close = rewards.tolerance(

distances, bounds=(0, _CLOSE_THRESHOLD), margin=(_CLOSE_THRESHOLD * 10))

# Fine-tolerance component for clicking the bricks together.

clicked = rewards.tolerance(

distances, bounds=(0, _CLICK_THRESHOLD), margin=_CLICK_THRESHOLD)

# Weighted average of coarse and fine components for each pair of bricks.

return np.average([close, clicked], weights=[close_coef, 1.], axis=0)

def _build_stack(physics, bricks, base_pos, base_quat, order, random_state):

"""Builds a stack of bricks.

Args:

physics: Instance of `mjcf.Physics`.

bricks: Sequence of `composer.Entity` instances corresponding to bricks.

base_pos: Position of the base brick in the stack; either a (3,) numpy array

or a `variation.Variation` that yields such arrays.

base_quat: Quaternion of the base brick in the stack; either a (4,) numpy

array or a `variation.Variation` that yields such arrays.

order: Sequence of indices specifying the order in which to stack the

bricks.

random_state: An `np.random.RandomState` instance.

"""

base_pos = variation.evaluate(base_pos, random_state=random_state)

base_quat = variation.evaluate(base_quat, random_state=random_state)

bricks[order[0]].set_pose(physics, position=base_pos, quaternion=base_quat)

for bottom_idx, top_idx in zip(order[:-1], order[1:]):

bottom = bricks[bottom_idx]

top = bricks[top_idx]

stud_pos = physics.bind(bottom.studs[0, 0]).xpos

_, quat = bottom.get_pose(physics)

# The reward function treats top left -> top left and top left -> bottom

# right configurations as identical, so the orientations of the bricks are

# randomized so that 50% of the time the top brick is rotated 180 degrees

# relative to the brick below.

if random_state.rand() < 0.5:

quat = quat.copy()

axis = np.array([0., 0., 1.])

angle = np.pi

mjlib.mju_quatIntegrate(quat, axis, angle)

hole_idx = (-1, -1)

else:

hole_idx = (0, 0)

top.set_pose(physics, quaternion=quat)

# Set the position of the top brick so that its holes line up with the studs

# of the brick below.

offset = physics.bind(top.holes[hole_idx]).xpos

top_pos = stud_pos - offset

top.set_pose(physics, position=top_pos)

def _add_or_remove_freejoints(attachment_frames, fixed_indices):

"""Adds or removes freejoints from props.

Args:

attachment_frames: A list of `mjcf.Elements` corresponding to attachment

frames.

fixed_indices: A list of indices of attachment frames that should be fixed

to the world (i.e. have their freejoints removed). Freejoints will be

added to all other elements in `attachment_frames` if they do not already

possess them.

"""

for i, frame in enumerate(attachment_frames):

if i in fixed_indices:

if frame.freejoint:

frame.freejoint.remove()

elif not frame.freejoint:

frame.add('freejoint')

def _replace_alpha(rgba, alpha=0.3):

new_rgba = rgba.copy()

new_rgba[3] = alpha

return new_rgba

def _hintify(entity, alpha=None):

"""Modifies an entity for use as a 'visual hint'.

Contacts will be disabled for all geoms within the entity, and its bodies will

be converted to "mocap" bodies (which are viewed as fixed from the perspective

of the dynamics). The geom alpha values may also be overridden to render the

geoms as translucent.

Args:

entity: A `composer.Entity`, modified in place.

alpha: Optional float between 0 and 1, used to override the alpha values for

all of the geoms in this entity.

"""

for subentity in entity.iter_entities():

# TODO(b/112084359): This assumes that all geoms either define explicit RGBA

# values, or inherit from the top-level default. It will

# not correctly handle more complicated hierarchies of

# default classes.

if (alpha is not None

and subentity.mjcf_model.default.geom is not None

and subentity.mjcf_model.default.geom.rgba is not None):

subentity.mjcf_model.default.geom.rgba = _replace_alpha(

subentity.mjcf_model.default.geom.rgba, alpha=alpha)

for body in subentity.mjcf_model.find_all('body'):

body.mocap = 'true'

for geom in subentity.mjcf_model.find_all('geom'):

if alpha is not None and geom.rgba is not None:

geom.rgba = _replace_alpha(geom.rgba, alpha=alpha)

geom.contype = 0

geom.conaffinity = 0

def _stack(obs_settings, num_bricks, moveable_base, randomize_order,

target_height=None):

"""Configure and instantiate a Stack task.

Args:

obs_settings: `observations.ObservationSettings` instance.

num_bricks: The total number of bricks; must be between 2 and 6.

moveable_base: Boolean specifying whether or not the bottom brick should

be moveable.

randomize_order: Boolean specifying whether to randomize the desired order

of bricks in the stack at the start of each episode.

target_height: The target number of bricks in the stack in order to get

maximum reward. Must be between 2 and `num_bricks`. Defaults to

`num_bricks`.

Returns:

An instance of `Stack`.

"""

if target_height is None:

target_height = num_bricks

arena = arenas.Standard()

arm = robots.make_arm(obs_settings=obs_settings)

hand = robots.make_hand(obs_settings=obs_settings)

return Stack(arena=arena,

arm=arm,

hand=hand,

num_bricks=num_bricks,

target_height=target_height,

moveable_base=moveable_base,

randomize_order=randomize_order,

obs_settings=obs_settings,

workspace=_WORKSPACE,

control_timestep=constants.CONTROL_TIMESTEP)

@registry.add(tags.FEATURES)

def stack_2_bricks_features():

return _stack(obs_settings=observations.PERFECT_FEATURES, num_bricks=2,

moveable_base=False, randomize_order=False)

@registry.add(tags.VISION)

def stack_2_bricks_vision():

return _stack(obs_settings=observations.VISION, num_bricks=2,

moveable_base=False, randomize_order=False)

@registry.add(tags.FEATURES)

def stack_2_bricks_moveable_base_features():

return _stack(obs_settings=observations.PERFECT_FEATURES, num_bricks=2,

moveable_base=True, randomize_order=False)

@registry.add(tags.VISION)

def stack_2_bricks_moveable_base_vision():

return _stack(obs_settings=observations.VISION, num_bricks=2,

moveable_base=True, randomize_order=False)

@registry.add(tags.FEATURES)

def stack_3_bricks_features():

return _stack(obs_settings=observations.PERFECT_FEATURES, num_bricks=3,

moveable_base=False, randomize_order=False)

@registry.add(tags.VISION)

def stack_3_bricks_vision():

return _stack(obs_settings=observations.VISION, num_bricks=3,

moveable_base=False, randomize_order=False)

@registry.add(tags.FEATURES)

def stack_3_bricks_random_order_features():

return _stack(obs_settings=observations.PERFECT_FEATURES, num_bricks=3,

moveable_base=False, randomize_order=True)

@registry.add(tags.FEATURES)

def stack_2_of_3_bricks_random_order_features():

return _stack(obs_settings=observations.PERFECT_FEATURES, num_bricks=3,

moveable_base=False, randomize_order=True, target_height=2)

@registry.add(tags.VISION)

def stack_2_of_3_bricks_random_order_vision():

return _stack(obs_settings=observations.VISION, num_bricks=3,

moveable_base=False, randomize_order=True, target_height=2)

def _reassemble(obs_settings, num_bricks, randomize_initial_order,

randomize_desired_order):

"""Configure and instantiate a `Reassemble` task.

Args:

obs_settings: `observations.ObservationSettings` instance.

num_bricks: The total number of bricks; must be between 2 and 6.

randomize_initial_order: Boolean specifying whether to randomize the

initial order of bricks in the stack at the start of each episode.

randomize_desired_order: Boolean specifying whether to independently

randomize the desired order of bricks in the stack at the start of each

episode. By default the desired order will be the reverse of the initial

order, with the exception of the base brick which is always the same as

in the initial order since it is welded in place.

Returns:

An instance of `Reassemble`.

"""

arena = arenas.Standard()

arm = robots.make_arm(obs_settings=obs_settings)

hand = robots.make_hand(obs_settings=obs_settings)

return Reassemble(arena=arena,

arm=arm,

hand=hand,

num_bricks=num_bricks,

randomize_initial_order=randomize_initial_order,

randomize_desired_order=randomize_desired_order,

obs_settings=obs_settings,

workspace=_WORKSPACE,

control_timestep=constants.CONTROL_TIMESTEP)

@registry.add(tags.FEATURES)

def reassemble_3_bricks_fixed_order_features():

return _reassemble(obs_settings=observations.PERFECT_FEATURES, num_bricks=3,

randomize_initial_order=False,

randomize_desired_order=False)

@registry.add(tags.VISION)

def reassemble_3_bricks_fixed_order_vision():

return _reassemble(obs_settings=observations.VISION, num_bricks=3,

randomize_initial_order=False,

randomize_desired_order=False)

@registry.add(tags.FEATURES)

def reassemble_5_bricks_random_order_features():

return _reassemble(obs_settings=observations.PERFECT_FEATURES, num_bricks=5,

randomize_initial_order=True,

randomize_desired_order=True)

@registry.add(tags.VISION)

def reassemble_5_bricks_random_order_vision():

return _reassemble(obs_settings=observations.VISION, num_bricks=5,

randomize_initial_order=True,

randomize_desired_order=True)