Sets the state of the environment at the start of each episode.
(self, physics)
| 184 | super().__init__(random=random) |
| 185 | |
| 186 | def initialize_episode(self, physics): |
| 187 | """Sets the state of the environment at the start of each episode.""" |
| 188 | # Local aliases |
| 189 | choice = self.random.choice |
| 190 | uniform = self.random.uniform |
| 191 | model = physics.named.model |
| 192 | data = physics.named.data |
| 193 | |
| 194 | # Find a collision-free random initial configuration. |
| 195 | penetrating = True |
| 196 | while penetrating: |
| 197 | |
| 198 | # Randomise angles of arm joints. |
| 199 | is_limited = model.jnt_limited[_ARM_JOINTS].astype(bool) |
| 200 | joint_range = model.jnt_range[_ARM_JOINTS] |
| 201 | lower_limits = np.where(is_limited, joint_range[:, 0], -np.pi) |
| 202 | upper_limits = np.where(is_limited, joint_range[:, 1], np.pi) |
| 203 | angles = uniform(lower_limits, upper_limits) |
| 204 | data.qpos[_ARM_JOINTS] = angles |
| 205 | |
| 206 | # Symmetrize hand. |
| 207 | data.qpos['finger'] = data.qpos['thumb'] |
| 208 | |
| 209 | # Randomise target location. |
| 210 | target_x = uniform(-.4, .4) |
| 211 | target_z = uniform(.1, .4) |
| 212 | if self._insert: |
| 213 | target_angle = uniform(-np.pi/3, np.pi/3) |
| 214 | model.body_pos[self._receptacle, ['x', 'z']] = target_x, target_z |
| 215 | model.body_quat[self._receptacle, ['qw', 'qy']] = [ |
| 216 | np.cos(target_angle/2), np.sin(target_angle/2)] |
| 217 | else: |
| 218 | target_angle = uniform(-np.pi, np.pi) |
| 219 | |
| 220 | model.body_pos[self._target, ['x', 'z']] = target_x, target_z |
| 221 | model.body_quat[self._target, ['qw', 'qy']] = [ |
| 222 | np.cos(target_angle/2), np.sin(target_angle/2)] |
| 223 | |
| 224 | # Randomise object location. |
| 225 | object_init_probs = [_P_IN_HAND, _P_IN_TARGET, 1-_P_IN_HAND-_P_IN_TARGET] |
| 226 | init_type = choice(['in_hand', 'in_target', 'uniform'], |
| 227 | p=object_init_probs) |
| 228 | if init_type == 'in_target': |
| 229 | object_x = target_x |
| 230 | object_z = target_z |
| 231 | object_angle = target_angle |
| 232 | elif init_type == 'in_hand': |
| 233 | physics.after_reset() |
| 234 | object_x = data.site_xpos['grasp', 'x'] |
| 235 | object_z = data.site_xpos['grasp', 'z'] |
| 236 | grasp_direction = data.site_xmat['grasp', ['xx', 'zx']] |
| 237 | object_angle = np.pi-np.arctan2(grasp_direction[1], grasp_direction[0]) |
| 238 | else: |
| 239 | object_x = uniform(-.5, .5) |
| 240 | object_z = uniform(0, .7) |
| 241 | object_angle = uniform(0, 2*np.pi) |
| 242 | data.qvel[self._object + '_x'] = uniform(-5, 5) |
| 243 |
nothing calls this directly
no test coverage detected