(self, physics)
Sets the state of the environment at the start of each episode.
Sets the state of the environment at the start of each episode.
Sets the state of the environment at the start of each episode .
def initialize_episode(self, physics):
    """Sets the state of the environment at the start of each episode."""
    # Local aliases
    choice = self.random.choice
    uniform = self.random.uniform
    model = physics.named.model
    data = physics.named.data

    # Find a collision-free random initial configuration.
    penetrating = True
    while penetrating:

      # Randomise angles of arm joints.
      is_limited = model.jnt_limited[_ARM_JOINTS].astype(np.bool)
      joint_range = model.jnt_range[_ARM_JOINTS]
      lower_limits = np.where(is_limited, joint_range[:, 0], -np.pi)
      upper_limits = np.where(is_limited, joint_range[:, 1], np.pi)
      angles = uniform(lower_limits, upper_limits)
      data.qpos[_ARM_JOINTS] = angles

      # Symmetrize hand.
      data.qpos['finger'] = data.qpos['thumb']

      # Randomise target location.
      target_x = uniform(-.4, .4)
      target_z = uniform(.1, .4)
      if self._insert:
        target_angle = uniform(-np.pi/3, np.pi/3)
        model.body_pos[self._receptacle, ['x', 'z']] = target_x, target_z
        model.body_quat[self._receptacle, ['qw', 'qy']] = [
            np.cos(target_angle/2), np.sin(target_angle/2)]
      else:
        target_angle = uniform(-np.pi, np.pi)

      model.body_pos[self._target, ['x', 'z']] = target_x, target_z
      model.body_quat[self._target, ['qw', 'qy']] = [
          np.cos(target_angle/2), np.sin(target_angle/2)]

      # Randomise object location.
      object_init_probs = [_P_IN_HAND, _P_IN_TARGET, 1-_P_IN_HAND-_P_IN_TARGET]
      init_type = choice(['in_hand', 'in_target', 'uniform'],
                         p=object_init_probs)
      if init_type == 'in_target':
        object_x = target_x
        object_z = target_z
        object_angle = target_angle
      elif init_type == 'in_hand':
        physics.after_reset()
        object_x = data.site_xpos['grasp', 'x']
        object_z = data.site_xpos['grasp', 'z']
        grasp_direction = data.site_xmat['grasp', ['xx', 'zx']]
        object_angle = np.pi-np.arctan2(grasp_direction[1], grasp_direction[0])
      else:
        object_x = uniform(-.5, .5)
        object_z = uniform(0, .7)
        object_angle = uniform(0, 2*np.pi)
        data.qvel[self._object + '_x'] = uniform(-5, 5)

      data.qpos[self._object_joints] = object_x, object_z, object_angle

      # Check for collisions.
      physics.after_reset()
      penetrating = physics.data.ncon > 0

    super().initialize_episode(physics)
def initialize_episode ( self , physics ) : # Local aliases choice = self . random . choice uniform = self . random . uniform model = physics . named . model data = physics . named . data # Find a collision-free random initial configuration. penetrating = True while penetrating : # Randomise angles of arm joints. is_limited = model . jnt_limited [ _ARM_JOINTS ] . astype ( np . bool ) joint_range = model . jnt_range [ _ARM_JOINTS ] lower_limits = np . where ( is_limited , joint_range [ : , 0 ] , - np . pi ) upper_limits = np . where ( is_limited , joint_range [ : , 1 ] , np . pi ) angles = uniform ( lower_limits , upper_limits ) data . qpos [ _ARM_JOINTS ] = angles # Symmetrize hand. data . qpos [ 'finger' ] = data . qpos [ 'thumb' ] # Randomise target location. target_x = uniform ( - .4 , .4 ) target_z = uniform ( .1 , .4 ) if self . _insert : target_angle = uniform ( - np . pi / 3 , np . pi / 3 ) model . body_pos [ self . _receptacle , [ 'x' , 'z' ] ] = target_x , target_z model . body_quat [ self . _receptacle , [ 'qw' , 'qy' ] ] = [ np . cos ( target_angle / 2 ) , np . sin ( target_angle / 2 ) ] else : target_angle = uniform ( - np . pi , np . pi ) model . body_pos [ self . _target , [ 'x' , 'z' ] ] = target_x , target_z model . body_quat [ self . _target , [ 'qw' , 'qy' ] ] = [ np . cos ( target_angle / 2 ) , np . sin ( target_angle / 2 ) ] # Randomise object location. object_init_probs = [ _P_IN_HAND , _P_IN_TARGET , 1 - _P_IN_HAND - _P_IN_TARGET ] init_type = choice ( [ 'in_hand' , 'in_target' , 'uniform' ] , p = object_init_probs ) if init_type == 'in_target' : object_x = target_x object_z = target_z object_angle = target_angle elif init_type == 'in_hand' : physics . after_reset ( ) object_x = data . site_xpos [ 'grasp' , 'x' ] object_z = data . site_xpos [ 'grasp' , 'z' ] grasp_direction = data . site_xmat [ 'grasp' , [ 'xx' , 'zx' ] ] object_angle = np . pi - np . arctan2 ( grasp_direction [ 1 ] , grasp_direction [ 0 ] ) else : object_x = uniform ( - .5 , .5 ) object_z = uniform ( 0 , .7 ) object_angle = uniform ( 0 , 2 * np . pi ) data . qvel [ self . _object + '_x' ] = uniform ( - 5 , 5 ) data . qpos [ self . _object_joints ] = object_x , object_z , object_angle # Check for collisions. physics . after_reset ( ) penetrating = physics . data . ncon > 0 super ( ) . initialize_episode ( physics )