A Stack `Task`: stack the boxes.
| 114 | |
| 115 | |
| 116 | class Stack(base.Task): |
| 117 | """A Stack `Task`: stack the boxes.""" |
| 118 | |
| 119 | def __init__(self, n_boxes, fully_observable, random=None): |
| 120 | """Initialize an instance of the `Stack` task. |
| 121 | |
| 122 | Args: |
| 123 | n_boxes: An `int`, number of boxes to stack. |
| 124 | fully_observable: A `bool`, whether the observation should contain the |
| 125 | positions and velocities of the boxes and the location of the target. |
| 126 | random: Optional, either a `numpy.random.RandomState` instance, an |
| 127 | integer seed for creating a new `RandomState`, or None to select a seed |
| 128 | automatically (default). |
| 129 | """ |
| 130 | self._n_boxes = n_boxes |
| 131 | self._box_names = ['box' + str(b) for b in range(n_boxes)] |
| 132 | self._box_joint_names = [] |
| 133 | for name in self._box_names: |
| 134 | for dim in 'xyz': |
| 135 | self._box_joint_names.append('_'.join([name, dim])) |
| 136 | self._fully_observable = fully_observable |
| 137 | super().__init__(random=random) |
| 138 | |
| 139 | def initialize_episode(self, physics): |
| 140 | """Sets the state of the environment at the start of each episode.""" |
| 141 | # Local aliases |
| 142 | randint = self.random.randint |
| 143 | uniform = self.random.uniform |
| 144 | model = physics.named.model |
| 145 | data = physics.named.data |
| 146 | |
| 147 | # Find a collision-free random initial configuration. |
| 148 | penetrating = True |
| 149 | while penetrating: |
| 150 | |
| 151 | # Randomise angles of arm joints. |
| 152 | is_limited = model.jnt_limited[_ARM_JOINTS].astype(bool) |
| 153 | joint_range = model.jnt_range[_ARM_JOINTS] |
| 154 | lower_limits = np.where(is_limited, joint_range[:, 0], -np.pi) |
| 155 | upper_limits = np.where(is_limited, joint_range[:, 1], np.pi) |
| 156 | angles = uniform(lower_limits, upper_limits) |
| 157 | data.qpos[_ARM_JOINTS] = angles |
| 158 | |
| 159 | # Symmetrize hand. |
| 160 | data.qpos['finger'] = data.qpos['thumb'] |
| 161 | |
| 162 | # Randomise target location. |
| 163 | target_height = 2*randint(self._n_boxes) + 1 |
| 164 | box_size = model.geom_size['target', 0] |
| 165 | model.body_pos['target', 'z'] = box_size * target_height |
| 166 | model.body_pos['target', 'x'] = uniform(-.37, .37) |
| 167 | |
| 168 | # Randomise box locations. |
| 169 | for name in self._box_names: |
| 170 | data.qpos[name + '_x'] = uniform(.1, .3) |
| 171 | data.qpos[name + '_z'] = uniform(0, .7) |
| 172 | data.qpos[name + '_y'] = uniform(0, 2*np.pi) |
| 173 |