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

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

"""A task where the goal is to place a movable prop on top of a fixed prop."""

import collections

from dm_control import composer

from dm_control import mjcf

from dm_control.composer import define

from dm_control.composer import initializers

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.utils import rewards

import numpy as np

_PlaceWorkspace = collections.namedtuple(

'_PlaceWorkspace', ['prop_bbox', 'target_bbox', 'tcp_bbox', 'arm_offset'])

_TARGET_RADIUS = 0.05

_PEDESTAL_RADIUS = 0.07

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

_PROP_Z_OFFSET = 1e-6

_WORKSPACE = _PlaceWorkspace(

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, _PEDESTAL_RADIUS + 0.1),

upper=(0.1, 0.1, 0.4)),

target_bbox=workspaces.BoundingBox(

lower=(-0.1, -0.1, _PEDESTAL_RADIUS),

upper=(0.1, 0.1, _PEDESTAL_RADIUS + 0.1)),

arm_offset=robots.ARM_OFFSET)

class SphereCradle(composer.Entity):

"""A concave shape for easy placement."""

_SPHERE_COUNT = 3

def _build(self):

self._mjcf_root = mjcf.element.RootElement(model='cradle')

sphere_radius = _PEDESTAL_RADIUS * 0.7

for ang in np.linspace(0, 2*np.pi, num=self._SPHERE_COUNT, endpoint=False):

pos = 0.7 * sphere_radius * np.array([np.sin(ang), np.cos(ang), -1])

self._mjcf_root.worldbody.add(

'geom', type='sphere', size=[sphere_radius], condim=6, pos=pos)

@property

def mjcf_model(self):

return self._mjcf_root

class Pedestal(composer.Entity):

"""A narrow pillar to elevate the target."""

_HEIGHT = 0.2

def _build(self, cradle, target_radius):

self._mjcf_root = mjcf.element.RootElement(model='pedestal')

self._mjcf_root.worldbody.add(

'geom', type='capsule', size=[_PEDESTAL_RADIUS],

fromto=[0, 0, -_PEDESTAL_RADIUS,

0, 0, -(self._HEIGHT + _PEDESTAL_RADIUS)])

attachment_site = self._mjcf_root.worldbody.add(

'site', type='sphere', size=(0.003,), group=constants.TASK_SITE_GROUP)

self.attach(cradle, attachment_site)

self._target_site = workspaces.add_target_site(

body=self.mjcf_model.worldbody,

radius=target_radius, rgba=constants.RED)

@property

def mjcf_model(self):

return self._mjcf_root

@property

def target_site(self):

return self._target_site

def _build_observables(self):

return PedestalObservables(self)

class PedestalObservables(composer.Observables):

"""Observables for the `Pedestal` prop."""

@define.observable

def position(self):

return observable.MJCFFeature('xpos', self._entity.target_site)

class Place(composer.Task):

"""Place the prop on top of another fixed prop held up by a pedestal."""

def __init__(self, arena, arm, hand, prop, obs_settings, workspace,

control_timestep, cradle):

"""Initializes a new `Place` task.

Args:

arena: `composer.Entity` instance.

arm: `robot_base.RobotArm` instance.

hand: `robot_base.RobotHand` instance.

prop: `composer.Entity` instance.

obs_settings: `observations.ObservationSettings` instance.

workspace: A `_PlaceWorkspace` instance.

control_timestep: Float specifying the control timestep in seconds.

cradle: `composer.Entity` onto which the `prop` must be placed.

"""

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)

self._tcp_initializer = initializers.ToolCenterPointInitializer(

self._hand, self._arm,

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

quaternion=workspaces.DOWN_QUATERNION)

self._prop = prop

self._prop_frame = self._arena.add_free_entity(prop)

self._pedestal = Pedestal(cradle=cradle, target_radius=_TARGET_RADIUS)

self._arena.attach(self._pedestal)

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

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

self._prop_placer = initializers.PropPlacer(

props=[prop],

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

quaternion=workspaces.uniform_z_rotation,

settle_physics=True,

max_attempts_per_prop=50)

self._pedestal_placer = initializers.PropPlacer(

props=[self._pedestal],

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

settle_physics=False)

# 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')

workspaces.add_bbox_site(

body=self.root_entity.mjcf_model.worldbody,

lower=workspace.target_bbox.lower,

upper=workspace.target_bbox.upper,

rgba=constants.CYAN, name='pedestal_spawn_area')

@property

def root_entity(self):

return self._arena

@property

def arm(self):

return self._arm

@property

def hand(self):

return self._hand

@property

def task_observables(self):

return self._task_observables

def initialize_episode(self, physics, random_state):

self._pedestal_placer(physics, random_state,

ignore_contacts_with_entities=[self._prop])

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

self._tcp_initializer(physics, random_state)

self._prop_placer(physics, random_state)

def get_reward(self, physics):

target = physics.bind(self._pedestal.target_site).xpos

obj = physics.bind(self._prop_frame).xpos

tcp = physics.bind(self._hand.tool_center_point).xpos

tcp_to_obj = np.linalg.norm(obj - tcp)

grasp = rewards.tolerance(tcp_to_obj,

bounds=(0, _TARGET_RADIUS),

margin=_TARGET_RADIUS,

sigmoid='long_tail')

obj_to_target = np.linalg.norm(obj - target)

in_place = rewards.tolerance(obj_to_target,

bounds=(0, _TARGET_RADIUS),

margin=_TARGET_RADIUS,

sigmoid='long_tail')

tcp_to_target = np.linalg.norm(tcp - target)

hand_away = rewards.tolerance(tcp_to_target,

bounds=(4*_TARGET_RADIUS, np.inf),

margin=3*_TARGET_RADIUS,

sigmoid='long_tail')

in_place_weight = 10.

grasp_or_hand_away = grasp * (1 - in_place) + hand_away * in_place

return (

grasp_or_hand_away + in_place_weight * in_place) / (1 + in_place_weight)

def _place(obs_settings, cradle_prop_name):

"""Configure and instantiate a Place task.

Args:

obs_settings: `observations.ObservationSettings` instance.

cradle_prop_name: The name of the prop onto which the Duplo brick must be

placed. Must be either 'duplo' or 'cradle'.

Returns:

An instance of `Place`.

Raises:

ValueError: If `prop_name` is neither 'duplo' nor 'cradle'.

"""

arena = arenas.Standard()

arm = robots.make_arm(obs_settings=obs_settings)

hand = robots.make_hand(obs_settings=obs_settings)

prop = props.Duplo(

observable_options=observations.make_options(

obs_settings, observations.FREEPROP_OBSERVABLES))

if cradle_prop_name == 'duplo':

cradle = props.Duplo()

elif cradle_prop_name == 'cradle':

cradle = SphereCradle()

else:

raise ValueError(

'`cradle_prop_name` must be either \'duplo\' or \'cradle\'.')

task = Place(arena=arena, arm=arm, hand=hand, prop=prop,

obs_settings=obs_settings,

workspace=_WORKSPACE,

control_timestep=constants.CONTROL_TIMESTEP,

cradle=cradle)

return task

@registry.add(tags.FEATURES)

def place_brick_features():

return _place(obs_settings=observations.PERFECT_FEATURES,

cradle_prop_name='duplo')

@registry.add(tags.VISION)

def place_brick_vision():

return _place(obs_settings=observations.VISION, cradle_prop_name='duplo')

@registry.add(tags.FEATURES)

def place_cradle_features():

return _place(obs_settings=observations.PERFECT_FEATURES,

cradle_prop_name='cradle')

@registry.add(tags.VISION)

def place_cradle_vision():

return _place(obs_settings=observations.VISION, cradle_prop_name='cradle')