def _get_features()

in tensorflow_datasets/robomimic/robomimic_ph/robomimic_ph.py [0:0]

• 74 lines of code
• 6 McCabe index (conditional complexity)

  def _get_features(self):
    obs_dim = _TASKS[self.builder_config.task]['object']
    states_dim = _TASKS[self.builder_config.task]['states']
    action_size = _TASKS[self.builder_config.task]['action_size']

    observation = {
        'object': _float_tensor_feature(obs_dim,),
        'robot0_eef_pos': _float_tensor_feature(3),
        'robot0_eef_quat': _float_tensor_feature(4),
        'robot0_eef_vel_ang': _float_tensor_feature(3),
        'robot0_eef_vel_lin': _float_tensor_feature(3),
        'robot0_gripper_qpos': _float_tensor_feature(2),
        'robot0_gripper_qvel': _float_tensor_feature(2),
        'robot0_joint_pos': _float_tensor_feature(7),
        'robot0_joint_pos_cos': _float_tensor_feature(7),
        'robot0_joint_pos_sin': _float_tensor_feature(7),
        'robot0_joint_vel': _float_tensor_feature(7),
    }
    if self.builder_config.task == 'transport':
      observation['robot1_eef_pos'] = _float_tensor_feature(3)
      observation['robot1_eef_quat'] = _float_tensor_feature(4)
      observation['robot1_eef_vel_ang'] = _float_tensor_feature(3)
      observation['robot1_eef_vel_lin'] = _float_tensor_feature(3)
      observation['robot1_gripper_qpos'] = _float_tensor_feature(2)
      observation['robot1_gripper_qvel'] = _float_tensor_feature(2)
      observation['robot1_joint_pos'] = _float_tensor_feature(7)
      observation['robot1_joint_pos_cos'] = _float_tensor_feature(7)
      observation['robot1_joint_pos_sin'] = _float_tensor_feature(7)
      observation['robot1_joint_vel'] = _float_tensor_feature(7)

    if 'image' in self.builder_config.filename:
      if self.builder_config.task == 'tool_hang':
        observation['robot0_eye_in_hand_image'] = _image_feature(240)
        observation['sideview_image'] = _image_feature(240)
      elif self.builder_config.task == 'transport':
        observation['robot0_eye_in_hand_image'] = _image_feature(84)
        observation['robot1_eye_in_hand_image'] = _image_feature(84)
        observation['shouldercamera0_image'] = _image_feature(84)
        observation['shouldercamera1_image'] = _image_feature(84)
      else:
        observation['agentview_image'] = _image_feature(84)
        observation['robot0_eye_in_hand_image'] = _image_feature(84)

    episode_metadata = {
        'train': tf.bool,
        'valid': tf.bool,
    }
    if self.builder_config.task != 'tool_hang':
      episode_metadata = {
          '20_percent': tf.bool,
          '20_percent_train': tf.bool,
          '20_percent_valid': tf.bool,
          '50_percent': tf.bool,
          '50_percent_train': tf.bool,
          '50_percent_valid': tf.bool,
          'train': tf.bool,
          'valid': tf.bool,
      }

    features = tfds.features.FeaturesDict({
        'horizon':
            tf.int32,
        'episode_id':
            tf.string,
        'steps':
            tfds.features.Dataset({
                'action': _float_tensor_feature(action_size),
                'observation': observation,
                'reward': tf.float64,
                'is_first': tf.bool,
                'is_last': tf.bool,
                'is_terminal': tf.bool,
                'discount': tf.int32,
                'states': _float_tensor_feature(states_dim),
            }),
        **episode_metadata,
    })
    return features