mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-01 22:19:48 +00:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| fdd9d92349 | |||
| 50192947fd | |||
| 99443a936d | |||
| 7a3298ea26 | |||
| 2d8f5f314e | |||
| 732a12108e | |||
| 29fcf057dc |
@@ -92,6 +92,20 @@ LIBERO supports two control modes — `relative` (default) and `absolute`. Diffe
|
||||
--env.control_mode=relative # or "absolute"
|
||||
```
|
||||
|
||||
### Reset performance
|
||||
|
||||
By default, LeRobot preserves LIBERO's hard-reset behavior. With fixed initial
|
||||
states enabled, you can opt into soft resets to skip rebuilding the simulator
|
||||
model and renderer on every episode:
|
||||
|
||||
```bash
|
||||
--env.init_states=true --env.hard_reset=false
|
||||
```
|
||||
|
||||
Soft resets are faster but are not bit-identical to hard resets after the
|
||||
environment's settling steps, so camera observations and policy results may
|
||||
differ slightly. Use hard resets when reproducing benchmark results.
|
||||
|
||||
### Policy inputs and outputs
|
||||
|
||||
**Observations:**
|
||||
|
||||
@@ -134,6 +134,20 @@ LIBERO-plus supports two control modes — `relative` (default) and `absolute`.
|
||||
--env.control_mode=relative # or "absolute"
|
||||
```
|
||||
|
||||
### Reset performance
|
||||
|
||||
By default, LeRobot preserves LIBERO's hard-reset behavior. With fixed initial
|
||||
states enabled, you can opt into soft resets to skip rebuilding the simulator
|
||||
model and renderer on every episode:
|
||||
|
||||
```bash
|
||||
--env.init_states=true --env.hard_reset=false
|
||||
```
|
||||
|
||||
Soft resets are faster but are not bit-identical to hard resets after the
|
||||
environment's settling steps, so camera observations and policy results may
|
||||
differ slightly. Use hard resets when reproducing benchmark results.
|
||||
|
||||
### Policy inputs and outputs
|
||||
|
||||
**Observations:**
|
||||
|
||||
@@ -328,6 +328,7 @@ class LiberoEnv(EnvConfig):
|
||||
render_mode: str = "rgb_array"
|
||||
camera_name: str = "agentview_image,robot0_eye_in_hand_image"
|
||||
init_states: bool = True
|
||||
hard_reset: bool = True
|
||||
camera_name_mapping: dict[str, str] | None = None
|
||||
observation_height: int = 360
|
||||
observation_width: int = 360
|
||||
@@ -356,6 +357,8 @@ class LiberoEnv(EnvConfig):
|
||||
def __post_init__(self):
|
||||
if self.fps <= 0:
|
||||
raise ValueError(f"fps must be positive, got {self.fps}")
|
||||
if not self.hard_reset and not self.init_states:
|
||||
raise ValueError("hard_reset=False requires init_states=True")
|
||||
|
||||
if self.obs_type == "pixels":
|
||||
self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
|
||||
@@ -416,6 +419,7 @@ class LiberoEnv(EnvConfig):
|
||||
"observation_height": self.observation_height,
|
||||
"observation_width": self.observation_width,
|
||||
"control_freq": self.fps,
|
||||
"hard_reset": self.hard_reset,
|
||||
}
|
||||
if self.task_ids is not None:
|
||||
kwargs["task_ids"] = self.task_ids
|
||||
|
||||
@@ -128,10 +128,13 @@ class LiberoEnv(gym.Env):
|
||||
control_freq: int = 20,
|
||||
control_mode: str = "relative",
|
||||
is_libero_plus: bool = False,
|
||||
hard_reset: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
if control_freq <= 0:
|
||||
raise ValueError(f"control_freq must be positive, got {control_freq}")
|
||||
if not hard_reset and not init_states:
|
||||
raise ValueError("hard_reset=False requires init_states=True")
|
||||
self.task_id = task_id
|
||||
self.is_libero_plus = is_libero_plus
|
||||
self.obs_type = obs_type
|
||||
@@ -158,6 +161,7 @@ class LiberoEnv(gym.Env):
|
||||
self.camera_name_mapping = camera_name_mapping
|
||||
self.num_steps_wait = num_steps_wait
|
||||
self.control_freq = control_freq
|
||||
self.hard_reset = hard_reset
|
||||
self.episode_index = episode_index
|
||||
self.episode_length = episode_length
|
||||
# Load once and keep
|
||||
@@ -265,6 +269,9 @@ class LiberoEnv(gym.Env):
|
||||
camera_heights=self.observation_height,
|
||||
camera_widths=self.observation_width,
|
||||
control_freq=self.control_freq,
|
||||
# Soft resets skip LIBERO's model and renderer rebuild. They are opt-in
|
||||
# because settle steps can make their observations differ from hard resets.
|
||||
hard_reset=self.hard_reset,
|
||||
)
|
||||
env.reset()
|
||||
self._env = env
|
||||
@@ -377,8 +384,9 @@ class LiberoEnv(gym.Env):
|
||||
}
|
||||
)
|
||||
observation = self._format_raw_obs(raw_obs)
|
||||
if terminated:
|
||||
self.reset()
|
||||
# Return the terminal observation unchanged. The caller owns resetting after
|
||||
# termination; vector envs created below use NEXT_STEP autoreset. Resetting here
|
||||
# would therefore reset twice and skip an initial state.
|
||||
truncated = False
|
||||
return observation, reward, terminated, truncated, info
|
||||
|
||||
@@ -476,6 +484,7 @@ def create_libero_envs(
|
||||
print(f"Restricting to task_ids={task_ids_filter}")
|
||||
|
||||
is_async = env_cls is gym.vector.AsyncVectorEnv
|
||||
is_sync = env_cls is gym.vector.SyncVectorEnv
|
||||
|
||||
out: dict[str, dict[int, Any]] = defaultdict(dict)
|
||||
for suite_name in suite_names:
|
||||
@@ -512,6 +521,10 @@ def create_libero_envs(
|
||||
cached_act_space = lazy.action_space
|
||||
cached_metadata = lazy.metadata
|
||||
out[suite_name][tid] = lazy
|
||||
elif is_sync:
|
||||
out[suite_name][tid] = gym.vector.SyncVectorEnv(
|
||||
fns, autoreset_mode=gym.vector.AutoresetMode.NEXT_STEP
|
||||
)
|
||||
else:
|
||||
out[suite_name][tid] = env_cls(fns)
|
||||
print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}")
|
||||
|
||||
@@ -212,7 +212,12 @@ class _LazyAsyncVectorEnv:
|
||||
|
||||
def _ensure(self) -> None:
|
||||
if self._env is None:
|
||||
self._env = gym.vector.AsyncVectorEnv(self._env_fns, context="forkserver", shared_memory=True)
|
||||
self._env = gym.vector.AsyncVectorEnv(
|
||||
self._env_fns,
|
||||
context="forkserver",
|
||||
shared_memory=True,
|
||||
autoreset_mode=gym.vector.AutoresetMode.NEXT_STEP,
|
||||
)
|
||||
|
||||
@property
|
||||
def unwrapped(self):
|
||||
|
||||
@@ -604,6 +604,12 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Optimized autoregressive decoding for FAST tokens using KV Caching.
|
||||
|
||||
Greedy decoding stops once every sequence emits the end-of-action marker. The
|
||||
returned tensor keeps its fixed shape, with positions not generated after the
|
||||
batch-wide stop left zero-filled. Stochastic decoding always runs to
|
||||
``max_decoding_steps`` so early stopping does not change the RNG state used by
|
||||
subsequent calls.
|
||||
"""
|
||||
if max_decoding_steps is None:
|
||||
max_decoding_steps = self.config.max_action_tokens
|
||||
@@ -612,6 +618,12 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
device = tokens.device
|
||||
lm_head = self.paligemma_with_expert.paligemma.lm_head
|
||||
|
||||
# detokenize_actions() cuts at the first "|", so greedy decoding can stop once
|
||||
# every sequence has emitted it. Keep stochastic decoding unchanged because
|
||||
# skipping multinomial calls would shift the RNG state for subsequent calls.
|
||||
end_of_action_token_id = self._paligemma_tokenizer.convert_tokens_to_ids("|")
|
||||
finished = torch.zeros(bsize, dtype=torch.bool, device=device) if temperature == 0 else None
|
||||
|
||||
# --- 1. PREFILL PHASE ---
|
||||
# Process Images + Text Prompt + BOS token once to populate the KV cache.
|
||||
|
||||
@@ -663,6 +675,10 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
# Initialize storage for generated tokens
|
||||
generated_action_tokens = torch.zeros((bsize, max_decoding_steps), dtype=torch.long, device=device)
|
||||
generated_action_tokens[:, 0] = next_token.squeeze(-1)
|
||||
if finished is not None:
|
||||
finished |= next_token.squeeze(-1) == end_of_action_token_id
|
||||
if bool(finished.all()):
|
||||
return generated_action_tokens
|
||||
|
||||
# Track valid tokens mask (0 for pad, 1 for valid)
|
||||
# We need this to tell the new token what it can attend to (images + text + past actions)
|
||||
@@ -713,6 +729,11 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
|
||||
generated_action_tokens[:, t] = next_token.squeeze(-1)
|
||||
|
||||
if finished is not None:
|
||||
finished |= next_token.squeeze(-1) == end_of_action_token_id
|
||||
if bool(finished.all()):
|
||||
break
|
||||
|
||||
return generated_action_tokens
|
||||
|
||||
|
||||
|
||||
@@ -77,11 +77,13 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
_command_when_disabled: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
|
||||
raw_observation = self.transition.get(TransitionKey.OBSERVATION)
|
||||
|
||||
if observation is None:
|
||||
if raw_observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
observation = raw_observation.copy()
|
||||
|
||||
if self.use_ik_solution and "IK_solution" in self.transition.get(TransitionKey.COMPLEMENTARY_DATA):
|
||||
q_raw = self.transition.get(TransitionKey.COMPLEMENTARY_DATA)["IK_solution"]
|
||||
else:
|
||||
@@ -311,10 +313,12 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
|
||||
)
|
||||
|
||||
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
|
||||
if observation is None:
|
||||
raw_observation = self.transition.get(TransitionKey.OBSERVATION)
|
||||
if raw_observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
observation = raw_observation.copy()
|
||||
|
||||
q_raw = np.array(
|
||||
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
|
||||
dtype=float,
|
||||
@@ -391,13 +395,15 @@ class GripperVelocityToJoint(RobotActionProcessorStep):
|
||||
discrete_gripper: bool = False
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
|
||||
raw_observation = self.transition.get(TransitionKey.OBSERVATION)
|
||||
|
||||
gripper_vel = action.pop("ee.gripper_vel")
|
||||
|
||||
if observation is None:
|
||||
if raw_observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
observation = raw_observation.copy()
|
||||
|
||||
q_raw = np.array(
|
||||
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
|
||||
dtype=float,
|
||||
@@ -583,10 +589,12 @@ class InverseKinematicsRLStep(ProcessorStep):
|
||||
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
|
||||
)
|
||||
|
||||
observation = new_transition.get(TransitionKey.OBSERVATION).copy()
|
||||
if observation is None:
|
||||
raw_observation = new_transition.get(TransitionKey.OBSERVATION)
|
||||
if raw_observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
observation = raw_observation.copy()
|
||||
|
||||
q_raw = np.array(
|
||||
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
|
||||
dtype=float,
|
||||
|
||||
@@ -240,6 +240,7 @@ Using JSON config file:
|
||||
|
||||
import abc
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
@@ -384,24 +385,35 @@ def _resolve_io_paths(
|
||||
return output_repo_id, input_path, output_path
|
||||
|
||||
|
||||
def _is_in_place(input_path: Path, output_path: Path) -> bool:
|
||||
"""Whether both paths point to the same dataset directory.
|
||||
|
||||
Uses os.path.samefile (device+inode) which is robust to case-insensitive filesystems, hardlinks
|
||||
and symlinks.
|
||||
"""
|
||||
try:
|
||||
return os.path.samefile(input_path, output_path)
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def get_output_path(
|
||||
repo_id: str,
|
||||
new_repo_id: str | None,
|
||||
root: Path | str | None,
|
||||
new_root: Path | str | None,
|
||||
) -> tuple[str, Path]:
|
||||
) -> tuple[str, Path, Path | None]:
|
||||
output_repo_id, input_path, output_path = _resolve_io_paths(repo_id, new_repo_id, root, new_root)
|
||||
|
||||
# In case of in-place modification, create a backup of the original dataset (if it exists)
|
||||
if output_path == input_path:
|
||||
# In case of in-place modification, create a backup of the original dataset (if it exists).
|
||||
backup_path: Path | None = None
|
||||
if _is_in_place(input_path, output_path):
|
||||
backup_path = input_path.with_name(input_path.name + "_old")
|
||||
|
||||
if input_path.exists():
|
||||
if backup_path.exists():
|
||||
shutil.rmtree(backup_path)
|
||||
shutil.move(input_path, backup_path)
|
||||
|
||||
return output_repo_id, output_path
|
||||
return output_repo_id, output_path, backup_path
|
||||
|
||||
|
||||
def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
||||
@@ -412,7 +424,7 @@ def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
||||
raise ValueError("episode_indices must be specified for delete_episodes operation")
|
||||
|
||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||
output_repo_id, output_dir = get_output_path(
|
||||
output_repo_id, output_dir, backup_path = get_output_path(
|
||||
cfg.repo_id,
|
||||
new_repo_id=cfg.new_repo_id,
|
||||
root=cfg.root,
|
||||
@@ -420,8 +432,8 @@ def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
||||
)
|
||||
|
||||
# In case of in-place modification, make the dataset point to the backup directory
|
||||
if output_dir == dataset.root:
|
||||
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
|
||||
if backup_path is not None:
|
||||
dataset.root = backup_path
|
||||
|
||||
logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
|
||||
new_dataset = delete_episodes(
|
||||
@@ -525,7 +537,7 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
|
||||
raise ValueError("feature_names must be specified for remove_feature operation")
|
||||
|
||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||
output_repo_id, output_dir = get_output_path(
|
||||
output_repo_id, output_dir, backup_path = get_output_path(
|
||||
cfg.repo_id,
|
||||
new_repo_id=cfg.new_repo_id,
|
||||
root=cfg.root,
|
||||
@@ -533,8 +545,8 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
|
||||
)
|
||||
|
||||
# In case of in-place modification, make the dataset point to the backup directory
|
||||
if output_dir == dataset.root:
|
||||
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
|
||||
if backup_path is not None:
|
||||
dataset.root = backup_path
|
||||
|
||||
logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
|
||||
new_dataset = remove_feature(
|
||||
@@ -671,7 +683,7 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
|
||||
cfg.new_root,
|
||||
default_new_repo_id=f"{cfg.repo_id}_recomputed_stats",
|
||||
)
|
||||
in_place = output_root == input_root
|
||||
in_place = _is_in_place(input_root, output_root)
|
||||
|
||||
if in_place and not cfg.operation.overwrite:
|
||||
raise ValueError(
|
||||
@@ -731,7 +743,7 @@ def handle_reencode_videos(cfg: EditDatasetConfig) -> None:
|
||||
cfg.new_root,
|
||||
default_new_repo_id=f"{cfg.repo_id}_reencoded",
|
||||
)
|
||||
in_place = output_root == input_root
|
||||
in_place = _is_in_place(input_root, output_root)
|
||||
|
||||
if in_place and not cfg.operation.overwrite:
|
||||
raise ValueError(
|
||||
|
||||
@@ -243,9 +243,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
# Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
|
||||
# Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
|
||||
force_cpu = cfg.trainable_config.device == "cpu"
|
||||
# Drive Accelerate's autocast from policy.dtype (bf16/fp16 activate it; float32/absent -> launcher default).
|
||||
# Drive Accelerate's autocast from policy.dtype (bf16/fp16 activate it; float32 -> full precision).
|
||||
has_policy_dtype = hasattr(cfg.trainable_config, "dtype")
|
||||
policy_dtype = getattr(cfg.trainable_config, "dtype", None)
|
||||
mixed_precision = {"bfloat16": "bf16", "float16": "fp16", "float32": "no"}.get(policy_dtype)
|
||||
# Policies without a `dtype` field fall back to `use_amp`, which would otherwise be
|
||||
# silently ignored here while lerobot-eval honors it. Follow torch.autocast's default
|
||||
# for the configured device so training and evaluation use the same precision.
|
||||
if not has_policy_dtype and getattr(cfg.trainable_config, "use_amp", False):
|
||||
device_type = torch.device(cfg.trainable_config.device).type
|
||||
autocast_dtype = torch.get_autocast_dtype(device_type)
|
||||
mixed_precision = {torch.bfloat16: "bf16", torch.float16: "fp16"}[autocast_dtype]
|
||||
accelerator = Accelerator(
|
||||
step_scheduler_with_optimizer=False,
|
||||
mixed_precision=mixed_precision,
|
||||
|
||||
@@ -17,9 +17,14 @@
|
||||
import pytest
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor.converters import create_transition
|
||||
from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
EEReferenceAndDelta,
|
||||
ForwardKinematicsJointsToEEAction,
|
||||
ForwardKinematicsJointsToEEObservation,
|
||||
GripperVelocityToJoint,
|
||||
InverseKinematicsEEToJoints,
|
||||
InverseKinematicsRLStep,
|
||||
)
|
||||
|
||||
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
|
||||
@@ -43,3 +48,38 @@ def test_fk_feature_schema(step_cls, bucket, feature_type):
|
||||
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
|
||||
assert set(out) == EE_KEYS
|
||||
assert {feature.type for feature in out.values()} == {feature_type}
|
||||
|
||||
|
||||
EE_ACTION = dict.fromkeys(EE_KEYS, 0.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("step", "action"),
|
||||
[
|
||||
(
|
||||
EEReferenceAndDelta(kinematics=None, end_effector_step_sizes={}, motor_names=MOTOR_NAMES),
|
||||
dict(EE_ACTION),
|
||||
),
|
||||
(InverseKinematicsEEToJoints(kinematics=None, motor_names=MOTOR_NAMES), dict(EE_ACTION)),
|
||||
(GripperVelocityToJoint(), {**EE_ACTION, "ee.gripper_vel": 0.0}),
|
||||
(InverseKinematicsRLStep(kinematics=None, motor_names=MOTOR_NAMES), dict(EE_ACTION)),
|
||||
],
|
||||
ids=[
|
||||
"ee_reference_and_delta",
|
||||
"inverse_kinematics_ee_to_joints",
|
||||
"gripper_velocity_to_joint",
|
||||
"inverse_kinematics_rl_step",
|
||||
],
|
||||
)
|
||||
def test_missing_observation_raises_value_error(step, action):
|
||||
"""A transition without an observation must surface the documented ValueError.
|
||||
|
||||
`RobotProcessorPipeline.process_action` builds its transition with
|
||||
`create_transition(action=...)`, which sets `TransitionKey.OBSERVATION` to None.
|
||||
These steps used to call `.copy()` on that before the None check, so the guard
|
||||
below them was unreachable and an AttributeError escaped instead.
|
||||
"""
|
||||
transition = create_transition(action=action)
|
||||
|
||||
with pytest.raises(ValueError, match="Joints observation"):
|
||||
step(transition)
|
||||
|
||||
Reference in New Issue
Block a user