google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Utilities for adroit environments."""
17

18
import enum
19

20

21
class VisibleState(enum.Enum):
22
  IMAGE = 1
23
  ROBOT = 2
24
  FULL = 3
25

26

27
def get_visible_features_for_task(task, visible_state):
28
  """Defines names of robot state or full state features for task."""
29
  features = []
30
  if isinstance(visible_state, str):
31
    try:
32
      visible_state = VisibleState[visible_state.upper()]
33
    except KeyError:
34
      # Custom subset of state features.
35
      features = visible_state.split(',')
36
  if visible_state == VisibleState.ROBOT:
37
    features = ['qpos', 'palm_pos', 'tactile']
38
    if task != 'relocate':
39
      features = ['qpos', 'qvel', 'palm_pos', 'tactile']
40
  elif visible_state == VisibleState.FULL:
41
    features = ['original_obs']
42
  return features
43