mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-01 03:11:29 +00:00
1e0d667a22
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com> Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>
105 lines
4.2 KiB
Python
105 lines
4.2 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import warnings
|
|
from typing import Any
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
from torch import Tensor
|
|
|
|
from lerobot.configs.types import FeatureType, PolicyFeature
|
|
from lerobot.envs.configs import EnvConfig
|
|
from lerobot.utils.utils import get_channel_first_image_shape
|
|
|
|
|
|
def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]:
|
|
"""Convert environment observation to LeRobot format observation.
|
|
|
|
This function uses the new pipeline system internally but maintains
|
|
backward compatibility with the original interface.
|
|
|
|
Args:
|
|
observation: Dictionary of observation batches from a Gym vector environment.
|
|
Returns:
|
|
Dictionary of observation batches with keys renamed to LeRobot format and values as tensors.
|
|
"""
|
|
from lerobot.processor import RobotProcessor, TransitionIndex, VanillaObservationProcessor
|
|
|
|
# Create processor with observation processor
|
|
processor = RobotProcessor([VanillaObservationProcessor()])
|
|
|
|
# Create transition tuple and process
|
|
transition = (observations, None, None, None, None, None, None)
|
|
processed_transition = processor(transition)
|
|
|
|
# Return processed observations
|
|
return processed_transition[TransitionIndex.OBSERVATION]
|
|
|
|
|
|
def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
|
|
# TODO(aliberts, rcadene): remove this hardcoding of keys and just use the nested keys as is
|
|
# (need to externalize normalization from policies)
|
|
policy_features = {}
|
|
for key, ft in env_cfg.features.items():
|
|
if ft.type is FeatureType.VISUAL:
|
|
if len(ft.shape) != 3:
|
|
raise ValueError(f"Number of dimensions of {key} != 3 (shape={ft.shape})")
|
|
|
|
shape = get_channel_first_image_shape(ft.shape)
|
|
feature = PolicyFeature(type=ft.type, shape=shape)
|
|
else:
|
|
feature = ft
|
|
|
|
policy_key = env_cfg.features_map[key]
|
|
policy_features[policy_key] = feature
|
|
|
|
return policy_features
|
|
|
|
|
|
def are_all_envs_same_type(env: gym.vector.VectorEnv) -> bool:
|
|
first_type = type(env.envs[0]) # Get type of first env
|
|
return all(type(e) is first_type for e in env.envs) # Fast type check
|
|
|
|
|
|
def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None:
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("once", UserWarning) # Apply filter only in this function
|
|
|
|
if not (hasattr(env.envs[0], "task_description") and hasattr(env.envs[0], "task")):
|
|
warnings.warn(
|
|
"The environment does not have 'task_description' and 'task'. Some policies require these features.",
|
|
UserWarning,
|
|
stacklevel=2,
|
|
)
|
|
if not are_all_envs_same_type(env):
|
|
warnings.warn(
|
|
"The environments have different types. Make sure you infer the right task from each environment. Empty task will be passed instead.",
|
|
UserWarning,
|
|
stacklevel=2,
|
|
)
|
|
|
|
|
|
def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]:
|
|
"""Adds task feature to the observation dict with respect to the first environment attribute."""
|
|
if hasattr(env.envs[0], "task_description"):
|
|
observation["task"] = env.call("task_description")
|
|
elif hasattr(env.envs[0], "task"):
|
|
observation["task"] = env.call("task")
|
|
else: # For envs without language instructions, e.g. aloha transfer cube and etc.
|
|
num_envs = observation[list(observation.keys())[0]].shape[0]
|
|
observation["task"] = ["" for _ in range(num_envs)]
|
|
return observation
|