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