mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 09:46:00 +00:00
fix(evo1): move LIBERO padding into policy processors
This commit is contained in:
@@ -139,6 +139,8 @@ every finetuning flag.
|
|||||||
| `policy.n_action_steps` | `50` | Number of actions consumed from a sampled chunk |
|
| `policy.n_action_steps` | `50` | Number of actions consumed from a sampled chunk |
|
||||||
| `policy.max_state_dim` | `24` | State padding dimension |
|
| `policy.max_state_dim` | `24` | State padding dimension |
|
||||||
| `policy.max_action_dim` | `24` | Action padding dimension |
|
| `policy.max_action_dim` | `24` | Action padding dimension |
|
||||||
|
| `policy.postprocess_action_dim` | `null` | Optional action dimension returned after EVO1 postprocessing |
|
||||||
|
| `policy.binarize_gripper` | `false` | Binarizes the postprocessed gripper channel for LIBERO-style eval |
|
||||||
| `policy.task_field` | `task` | Batch field used as the language prompt |
|
| `policy.task_field` | `task` | Batch field used as the language prompt |
|
||||||
|
|
||||||
## Results
|
## Results
|
||||||
@@ -161,16 +163,20 @@ pixel embeddings, VLM fused tokens, normalized actions, and denormalized actions
|
|||||||
The published checkpoint expects the raw LIBERO camera feature names
|
The published checkpoint expects the raw LIBERO camera feature names
|
||||||
`observation.images.agentview_image` and `observation.images.robot0_eye_in_hand_image`. The official EVO1 LIBERO
|
`observation.images.agentview_image` and `observation.images.robot0_eye_in_hand_image`. The official EVO1 LIBERO
|
||||||
rollout protocol also replans every 14 actions and binarizes the gripper command before stepping the simulator.
|
rollout protocol also replans every 14 actions and binarizes the gripper command before stepping the simulator.
|
||||||
The LIBERO environment postprocessor applies the gripper binarization automatically for EVO1 policies. To run the
|
The EVO1 policy postprocessor can crop the padded 24D action back to the 7D LIBERO action space and apply that
|
||||||
converted checkpoint with LeRobot LIBERO evaluation for the same one-episode-per-task setting, keep the raw camera
|
gripper binarization. To run the converted checkpoint with LeRobot LIBERO evaluation for the same
|
||||||
names instead of the default `image`/`image2` mapping and override `policy.n_action_steps` to 14:
|
one-episode-per-task setting, keep the raw camera names instead of the default `image`/`image2` mapping, enable
|
||||||
|
FlashAttention, and set the LIBERO action postprocessing flags:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
lerobot-eval \
|
lerobot-eval \
|
||||||
--policy.path=javadcc/evo1-libero-lerobot \
|
--policy.path=javadcc/evo1-libero-lerobot \
|
||||||
--policy.vlm_model_name=OpenGVLab/InternVL3-1B \
|
--policy.vlm_model_name=OpenGVLab/InternVL3-1B \
|
||||||
--policy.device=cuda \
|
--policy.device=cuda \
|
||||||
|
--policy.use_flash_attn=true \
|
||||||
--policy.n_action_steps=14 \
|
--policy.n_action_steps=14 \
|
||||||
|
--policy.postprocess_action_dim=7 \
|
||||||
|
--policy.binarize_gripper=true \
|
||||||
--env.type=libero \
|
--env.type=libero \
|
||||||
--env.task=libero_object \
|
--env.task=libero_object \
|
||||||
--env.camera_name_mapping="{agentview_image: agentview_image, robot0_eye_in_hand_image: robot0_eye_in_hand_image}" \
|
--env.camera_name_mapping="{agentview_image: agentview_image, robot0_eye_in_hand_image: robot0_eye_in_hand_image}" \
|
||||||
|
|||||||
+3
-3
@@ -140,6 +140,7 @@ pyserial-dep = ["pyserial>=3.5,<4.0"]
|
|||||||
deepdiff-dep = ["deepdiff>=7.0.1,<9.0.0"]
|
deepdiff-dep = ["deepdiff>=7.0.1,<9.0.0"]
|
||||||
pynput-dep = ["pynput>=1.7.8,<1.9.0"]
|
pynput-dep = ["pynput>=1.7.8,<1.9.0"]
|
||||||
pyzmq-dep = ["pyzmq>=26.2.1,<28.0.0"]
|
pyzmq-dep = ["pyzmq>=26.2.1,<28.0.0"]
|
||||||
|
timm-dep = ["timm>=1.0.0,<1.1.0"]
|
||||||
|
|
||||||
# Motors
|
# Motors
|
||||||
feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0", "lerobot[pyserial-dep]", "lerobot[deepdiff-dep]"]
|
feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0", "lerobot[pyserial-dep]", "lerobot[deepdiff-dep]"]
|
||||||
@@ -187,7 +188,7 @@ groot = [
|
|||||||
"lerobot[peft-dep]",
|
"lerobot[peft-dep]",
|
||||||
"lerobot[diffusers-dep]",
|
"lerobot[diffusers-dep]",
|
||||||
"dm-tree>=0.1.8,<1.0.0",
|
"dm-tree>=0.1.8,<1.0.0",
|
||||||
"timm>=1.0.0,<1.1.0",
|
"lerobot[timm-dep]",
|
||||||
"decord>=0.6.0,<1.0.0; (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
|
"decord>=0.6.0,<1.0.0; (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
|
||||||
"ninja>=1.11.1,<2.0.0",
|
"ninja>=1.11.1,<2.0.0",
|
||||||
"flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'"
|
"flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'"
|
||||||
@@ -195,7 +196,7 @@ groot = [
|
|||||||
sarm = ["lerobot[transformers-dep]", "pydantic>=2.0.0,<3.0.0", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"]
|
sarm = ["lerobot[transformers-dep]", "pydantic>=2.0.0,<3.0.0", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"]
|
||||||
xvla = ["lerobot[transformers-dep]"]
|
xvla = ["lerobot[transformers-dep]"]
|
||||||
eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"]
|
eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"]
|
||||||
evo1 = ["lerobot[transformers-dep]", "timm>=1.0.0,<1.1.0"]
|
evo1 = ["lerobot[transformers-dep]", "lerobot[timm-dep]"]
|
||||||
hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
|
hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
|
||||||
|
|
||||||
# Features
|
# Features
|
||||||
@@ -350,7 +351,6 @@ ignore = [
|
|||||||
# E402: conditional-import guards (TYPE_CHECKING / is_package_available) must precede the imports they protect
|
# E402: conditional-import guards (TYPE_CHECKING / is_package_available) must precede the imports they protect
|
||||||
"src/lerobot/scripts/convert_dataset_v21_to_v30.py" = ["E402"]
|
"src/lerobot/scripts/convert_dataset_v21_to_v30.py" = ["E402"]
|
||||||
"src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
|
"src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
|
||||||
"src/lerobot/policies/evo1/**" = ["N801", "N812"]
|
|
||||||
|
|
||||||
[tool.ruff.lint.isort]
|
[tool.ruff.lint.isort]
|
||||||
combine-as-imports = true
|
combine-as-imports = true
|
||||||
|
|||||||
@@ -26,7 +26,6 @@ from gymnasium.envs.registration import registry as gym_registry
|
|||||||
from lerobot.configs import FeatureType, PolicyFeature
|
from lerobot.configs import FeatureType, PolicyFeature
|
||||||
from lerobot.processor import (
|
from lerobot.processor import (
|
||||||
IsaaclabArenaProcessorStep,
|
IsaaclabArenaProcessorStep,
|
||||||
LiberoActionProcessorStep,
|
|
||||||
LiberoProcessorStep,
|
LiberoProcessorStep,
|
||||||
PolicyProcessorPipeline,
|
PolicyProcessorPipeline,
|
||||||
)
|
)
|
||||||
@@ -128,7 +127,7 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
|
|||||||
vec = env_cls([_make_one for _ in range(n_envs)], **extra_kwargs)
|
vec = env_cls([_make_one for _ in range(n_envs)], **extra_kwargs)
|
||||||
return {self.type: {0: vec}}
|
return {self.type: {0: vec}}
|
||||||
|
|
||||||
def get_env_processors(self, policy_cfg: Any | None = None):
|
def get_env_processors(self):
|
||||||
"""Return (preprocessor, postprocessor) for this env. Default: identity."""
|
"""Return (preprocessor, postprocessor) for this env. Default: identity."""
|
||||||
return PolicyProcessorPipeline(steps=[]), PolicyProcessorPipeline(steps=[])
|
return PolicyProcessorPipeline(steps=[]), PolicyProcessorPipeline(steps=[])
|
||||||
|
|
||||||
@@ -357,7 +356,6 @@ class LiberoEnv(EnvConfig):
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
control_mode: str = "relative" # or "absolute"
|
control_mode: str = "relative" # or "absolute"
|
||||||
binarize_gripper: bool | None = None
|
|
||||||
|
|
||||||
def __post_init__(self):
|
def __post_init__(self):
|
||||||
if self.obs_type == "pixels":
|
if self.obs_type == "pixels":
|
||||||
@@ -442,22 +440,10 @@ class LiberoEnv(EnvConfig):
|
|||||||
is_libero_plus=self.is_libero_plus,
|
is_libero_plus=self.is_libero_plus,
|
||||||
)
|
)
|
||||||
|
|
||||||
def get_env_processors(self, policy_cfg: Any | None = None):
|
def get_env_processors(self):
|
||||||
is_evo1 = getattr(policy_cfg, "type", None) == "evo1"
|
|
||||||
max_state_dim = getattr(policy_cfg, "max_state_dim", None) if is_evo1 else None
|
|
||||||
action_feature = self.features.get(ACTION)
|
|
||||||
action_dim = int(action_feature.shape[0]) if action_feature is not None else 7
|
|
||||||
binarize_gripper = is_evo1 if self.binarize_gripper is None else self.binarize_gripper
|
|
||||||
return (
|
return (
|
||||||
PolicyProcessorPipeline(steps=[LiberoProcessorStep(max_state_dim=max_state_dim)]),
|
PolicyProcessorPipeline(steps=[LiberoProcessorStep()]),
|
||||||
PolicyProcessorPipeline(
|
PolicyProcessorPipeline(steps=[]),
|
||||||
steps=[
|
|
||||||
LiberoActionProcessorStep(
|
|
||||||
action_dim=action_dim,
|
|
||||||
binarize_gripper=binarize_gripper,
|
|
||||||
)
|
|
||||||
]
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -723,7 +709,7 @@ class IsaaclabArenaEnv(HubEnvConfig):
|
|||||||
def gym_kwargs(self) -> dict:
|
def gym_kwargs(self) -> dict:
|
||||||
return {}
|
return {}
|
||||||
|
|
||||||
def get_env_processors(self, policy_cfg: Any | None = None):
|
def get_env_processors(self):
|
||||||
state_keys = tuple(k.strip() for k in (self.state_keys or "").split(",") if k.strip())
|
state_keys = tuple(k.strip() for k in (self.state_keys or "").split(",") if k.strip())
|
||||||
camera_keys = tuple(k.strip() for k in (self.camera_keys or "").split(",") if k.strip())
|
camera_keys = tuple(k.strip() for k in (self.camera_keys or "").split(",") if k.strip())
|
||||||
if not state_keys and not camera_keys:
|
if not state_keys and not camera_keys:
|
||||||
|
|||||||
@@ -15,7 +15,6 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import inspect
|
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import gymnasium as gym
|
import gymnasium as gym
|
||||||
@@ -53,14 +52,7 @@ def make_env_pre_post_processors(
|
|||||||
|
|
||||||
return make_xvla_libero_pre_post_processors()
|
return make_xvla_libero_pre_post_processors()
|
||||||
|
|
||||||
get_processors = env_cfg.get_env_processors
|
return env_cfg.get_env_processors()
|
||||||
signature = inspect.signature(get_processors)
|
|
||||||
supports_policy_cfg = "policy_cfg" in signature.parameters or any(
|
|
||||||
param.kind is inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()
|
|
||||||
)
|
|
||||||
if supports_policy_cfg:
|
|
||||||
return get_processors(policy_cfg=policy_cfg)
|
|
||||||
return get_processors()
|
|
||||||
|
|
||||||
|
|
||||||
def make_env(
|
def make_env(
|
||||||
|
|||||||
@@ -1 +0,0 @@
|
|||||||
../../../../docs/source/policy_evo1_README.md
|
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
# EVO1
|
||||||
|
|
||||||
|
EVO1 is a Vision-Language-Action policy for robot control. The LeRobot
|
||||||
|
integration uses an InternVL3 vision-language backbone with a flow-matching
|
||||||
|
action head, and supports staged training through the standard LeRobot policy
|
||||||
|
APIs.
|
||||||
|
|
||||||
|
The upstream EVO1 project is available at
|
||||||
|
[MINT-SJTU/Evo-1](https://github.com/MINT-SJTU/Evo-1).
|
||||||
|
|
||||||
|
```bibtex
|
||||||
|
@misc{evo1,
|
||||||
|
title = {EVO1},
|
||||||
|
author = {{MINT-SJTU}},
|
||||||
|
year = {2026},
|
||||||
|
howpublished = {\url{https://github.com/MINT-SJTU/Evo-1}},
|
||||||
|
}
|
||||||
|
```
|
||||||
@@ -14,6 +14,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
import math
|
import math
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
@@ -26,6 +27,8 @@ from lerobot.optim.optimizers import AdamWConfig
|
|||||||
from lerobot.optim.schedulers import LRSchedulerConfig
|
from lerobot.optim.schedulers import LRSchedulerConfig
|
||||||
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
|
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@LRSchedulerConfig.register_subclass("evo1_exact")
|
@LRSchedulerConfig.register_subclass("evo1_exact")
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -59,6 +62,12 @@ class Evo1Config(PreTrainedConfig):
|
|||||||
max_views: int = 3
|
max_views: int = 3
|
||||||
image_resolution: tuple[int, int] = (448, 448)
|
image_resolution: tuple[int, int] = (448, 448)
|
||||||
empty_cameras: int = 0
|
empty_cameras: int = 0
|
||||||
|
postprocess_action_dim: int | None = None
|
||||||
|
binarize_gripper: bool = False
|
||||||
|
gripper_index: int = 6
|
||||||
|
gripper_threshold: float = 0.5
|
||||||
|
gripper_below_threshold_value: float = 1.0
|
||||||
|
gripper_above_threshold_value: float = -1.0
|
||||||
|
|
||||||
normalization_mapping: dict[str, NormalizationMode] = field(
|
normalization_mapping: dict[str, NormalizationMode] = field(
|
||||||
default_factory=lambda: {
|
default_factory=lambda: {
|
||||||
@@ -114,16 +123,32 @@ class Evo1Config(PreTrainedConfig):
|
|||||||
)
|
)
|
||||||
|
|
||||||
if self.apply_training_stage_defaults:
|
if self.apply_training_stage_defaults:
|
||||||
if self.training_stage == "stage1":
|
stage_defaults = {
|
||||||
self.finetune_vlm = False
|
"stage1": {
|
||||||
self.finetune_language_model = False
|
"finetune_vlm": False,
|
||||||
self.finetune_vision_model = False
|
"finetune_language_model": False,
|
||||||
self.finetune_action_head = True
|
"finetune_vision_model": False,
|
||||||
elif self.training_stage == "stage2":
|
"finetune_action_head": True,
|
||||||
self.finetune_vlm = True
|
},
|
||||||
self.finetune_language_model = True
|
"stage2": {
|
||||||
self.finetune_vision_model = True
|
"finetune_vlm": True,
|
||||||
self.finetune_action_head = True
|
"finetune_language_model": True,
|
||||||
|
"finetune_vision_model": True,
|
||||||
|
"finetune_action_head": True,
|
||||||
|
},
|
||||||
|
}[self.training_stage]
|
||||||
|
for flag_name, default_value in stage_defaults.items():
|
||||||
|
current_value = getattr(self, flag_name)
|
||||||
|
if current_value is not None and current_value != default_value:
|
||||||
|
logger.warning(
|
||||||
|
"EVO1 %s=%s is overridden by training_stage=%s default %s. "
|
||||||
|
"Set apply_training_stage_defaults=false to keep explicit finetuning flags.",
|
||||||
|
flag_name,
|
||||||
|
current_value,
|
||||||
|
self.training_stage,
|
||||||
|
default_value,
|
||||||
|
)
|
||||||
|
setattr(self, flag_name, default_value)
|
||||||
elif self.training_stage == "stage1":
|
elif self.training_stage == "stage1":
|
||||||
if self.finetune_vlm is None:
|
if self.finetune_vlm is None:
|
||||||
self.finetune_vlm = False
|
self.finetune_vlm = False
|
||||||
@@ -171,6 +196,11 @@ class Evo1Config(PreTrainedConfig):
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"n_action_steps ({self.n_action_steps}) must be <= chunk_size ({self.chunk_size})"
|
f"n_action_steps ({self.n_action_steps}) must be <= chunk_size ({self.chunk_size})"
|
||||||
)
|
)
|
||||||
|
if len(self.image_resolution) != 2 or self.image_resolution[0] != self.image_resolution[1]:
|
||||||
|
raise ValueError(
|
||||||
|
"EVO1 currently expects a square image_resolution because InternVL3 preprocessing "
|
||||||
|
f"uses a scalar image_size, got {self.image_resolution}."
|
||||||
|
)
|
||||||
|
|
||||||
def validate_features(self) -> None:
|
def validate_features(self) -> None:
|
||||||
if self.input_features is None:
|
if self.input_features is None:
|
||||||
|
|||||||
@@ -21,8 +21,8 @@ import torch
|
|||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
from lerobot.policies.evo1.flow_matching import FlowmatchingActionHead
|
from .flow_matching import FlowmatchingActionHead
|
||||||
from lerobot.policies.evo1.internvl3_embedder import InternVL3Embedder
|
from .internvl3_embedder import InternVL3Embedder
|
||||||
|
|
||||||
|
|
||||||
def _cfgget(config: Any, key: str, default=None):
|
def _cfgget(config: Any, key: str, default=None):
|
||||||
@@ -163,37 +163,6 @@ class EVO1(nn.Module):
|
|||||||
embodiment_id=embodiment_ids,
|
embodiment_id=embodiment_ids,
|
||||||
)
|
)
|
||||||
|
|
||||||
@torch.no_grad()
|
|
||||||
def run_inference(
|
|
||||||
self,
|
|
||||||
images: list[Image.Image | torch.Tensor],
|
|
||||||
image_mask: torch.Tensor,
|
|
||||||
prompt: str,
|
|
||||||
state_input: list | torch.Tensor,
|
|
||||||
return_cls_only: bool | None = None,
|
|
||||||
action_mask: torch.Tensor | None = None,
|
|
||||||
embodiment_ids: torch.Tensor | None = None,
|
|
||||||
) -> torch.Tensor:
|
|
||||||
if image_mask.dim() == 1:
|
|
||||||
image_mask = image_mask.unsqueeze(0)
|
|
||||||
|
|
||||||
fused_tokens = self.get_vl_embeddings(
|
|
||||||
images=[images],
|
|
||||||
image_mask=image_mask,
|
|
||||||
prompt=[prompt],
|
|
||||||
return_cls_only=return_cls_only,
|
|
||||||
)
|
|
||||||
state_tensor = self.prepare_state(state_input)
|
|
||||||
action = self.predict_action(
|
|
||||||
fused_tokens,
|
|
||||||
state_tensor,
|
|
||||||
action_mask=action_mask,
|
|
||||||
embodiment_ids=embodiment_ids,
|
|
||||||
)
|
|
||||||
if isinstance(action, torch.Tensor) and action.dtype == torch.bfloat16:
|
|
||||||
action = action.to(torch.float32)
|
|
||||||
return action
|
|
||||||
|
|
||||||
def forward(
|
def forward(
|
||||||
self,
|
self,
|
||||||
fused_tokens: torch.Tensor,
|
fused_tokens: torch.Tensor,
|
||||||
|
|||||||
@@ -129,7 +129,10 @@ class MultiEmbodimentActionEncoder(nn.Module):
|
|||||||
|
|
||||||
def forward(self, action_seq: torch.Tensor, category_id: torch.LongTensor):
|
def forward(self, action_seq: torch.Tensor, category_id: torch.LongTensor):
|
||||||
batch_size, horizon, action_dim = action_seq.shape
|
batch_size, horizon, action_dim = action_seq.shape
|
||||||
assert self.horizon == horizon, "Action sequence length must match horizon"
|
if self.horizon != horizon:
|
||||||
|
raise ValueError(
|
||||||
|
f"Action sequence length must match horizon: got {horizon}, expected {self.horizon}."
|
||||||
|
)
|
||||||
|
|
||||||
x = action_seq.reshape(batch_size * horizon, action_dim)
|
x = action_seq.reshape(batch_size * horizon, action_dim)
|
||||||
if category_id.dim() == 0:
|
if category_id.dim() == 0:
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING
|
|||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
import torch.utils.checkpoint
|
import torch.utils.checkpoint
|
||||||
import torchvision.transforms.functional as TF
|
import torchvision.transforms.functional as tvf
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
from torchvision.transforms.functional import to_pil_image
|
from torchvision.transforms.functional import to_pil_image
|
||||||
|
|
||||||
@@ -46,6 +46,26 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
def _patch_vision_encoder_checkpointing(encoder: nn.Module, use_reentrant: bool) -> None:
|
def _patch_vision_encoder_checkpointing(encoder: nn.Module, use_reentrant: bool) -> None:
|
||||||
|
for attr_name in ("_gradient_checkpointing_func", "gradient_checkpointing_func"):
|
||||||
|
original_func = getattr(encoder, attr_name, None)
|
||||||
|
if not callable(original_func):
|
||||||
|
continue
|
||||||
|
patch_attr = f"_evo1_{attr_name}_patch_applied"
|
||||||
|
if getattr(encoder, patch_attr, False):
|
||||||
|
encoder.gradient_checkpointing_use_reentrant = use_reentrant
|
||||||
|
return
|
||||||
|
|
||||||
|
def checkpoint_with_kwargs(
|
||||||
|
function, *checkpoint_args, _original_func=original_func, **checkpoint_kwargs
|
||||||
|
):
|
||||||
|
checkpoint_kwargs.setdefault("use_reentrant", encoder.gradient_checkpointing_use_reentrant)
|
||||||
|
return _original_func(function, *checkpoint_args, **checkpoint_kwargs)
|
||||||
|
|
||||||
|
encoder.gradient_checkpointing_use_reentrant = use_reentrant
|
||||||
|
setattr(encoder, attr_name, checkpoint_with_kwargs)
|
||||||
|
setattr(encoder, patch_attr, True)
|
||||||
|
return
|
||||||
|
|
||||||
if getattr(encoder, "_evo1_checkpoint_patch_applied", False):
|
if getattr(encoder, "_evo1_checkpoint_patch_applied", False):
|
||||||
encoder.gradient_checkpointing_use_reentrant = use_reentrant
|
encoder.gradient_checkpointing_use_reentrant = use_reentrant
|
||||||
return
|
return
|
||||||
@@ -59,6 +79,9 @@ def _patch_vision_encoder_checkpointing(encoder: nn.Module, use_reentrant: bool)
|
|||||||
checkpoint_kwargs.setdefault("use_reentrant", self.gradient_checkpointing_use_reentrant)
|
checkpoint_kwargs.setdefault("use_reentrant", self.gradient_checkpointing_use_reentrant)
|
||||||
return original_checkpoint(function, *checkpoint_args, **checkpoint_kwargs)
|
return original_checkpoint(function, *checkpoint_args, **checkpoint_kwargs)
|
||||||
|
|
||||||
|
# Some InternVL3 remote-code versions call torch.utils.checkpoint.checkpoint
|
||||||
|
# directly and do not expose a per-encoder checkpoint function to patch.
|
||||||
|
# Keep this compatibility fallback scoped to encoder.forward and restore it.
|
||||||
torch.utils.checkpoint.checkpoint = checkpoint
|
torch.utils.checkpoint.checkpoint = checkpoint
|
||||||
try:
|
try:
|
||||||
return original_forward(*args, **kwargs)
|
return original_forward(*args, **kwargs)
|
||||||
@@ -280,11 +303,13 @@ class InternVL3Embedder(nn.Module):
|
|||||||
|
|
||||||
def _preprocess_single_image(self, image: Image.Image | torch.Tensor) -> torch.Tensor:
|
def _preprocess_single_image(self, image: Image.Image | torch.Tensor) -> torch.Tensor:
|
||||||
if isinstance(image, torch.Tensor):
|
if isinstance(image, torch.Tensor):
|
||||||
|
# Match upstream EVO1/InternVL preprocessing, which converts tensors
|
||||||
|
# through PIL before tiling and ImageNet normalization.
|
||||||
pil_image = to_pil_image(image.detach().cpu())
|
pil_image = to_pil_image(image.detach().cpu())
|
||||||
else:
|
else:
|
||||||
pil_image = image.convert("RGB")
|
pil_image = image.convert("RGB")
|
||||||
tiles = dynamic_preprocess(pil_image, image_size=self.image_size)
|
tiles = dynamic_preprocess(pil_image, image_size=self.image_size)
|
||||||
tile_tensors = torch.stack([TF.to_tensor(tile) for tile in tiles]).to(
|
tile_tensors = torch.stack([tvf.to_tensor(tile) for tile in tiles]).to(
|
||||||
device=self.device, dtype=torch.bfloat16
|
device=self.device, dtype=torch.bfloat16
|
||||||
)
|
)
|
||||||
mean = torch.tensor(IMAGENET_MEAN, device=self.device, dtype=torch.bfloat16).view(1, 3, 1, 1)
|
mean = torch.tensor(IMAGENET_MEAN, device=self.device, dtype=torch.bfloat16).view(1, 3, 1, 1)
|
||||||
|
|||||||
@@ -45,6 +45,7 @@ class EVO1Policy(PreTrainedPolicy):
|
|||||||
self.config = config
|
self.config = config
|
||||||
self.model = EVO1(self._build_model_config(config))
|
self.model = EVO1(self._build_model_config(config))
|
||||||
self.model.set_finetune_flags()
|
self.model.set_finetune_flags()
|
||||||
|
self._keep_frozen_embedder_eval()
|
||||||
self.reset()
|
self.reset()
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -64,7 +65,7 @@ class EVO1Policy(PreTrainedPolicy):
|
|||||||
**kwargs,
|
**kwargs,
|
||||||
) -> T:
|
) -> T:
|
||||||
if strict is None:
|
if strict is None:
|
||||||
strict = not (config is not None and getattr(config, "training_stage", None) == "stage2")
|
strict = True
|
||||||
return super().from_pretrained(
|
return super().from_pretrained(
|
||||||
pretrained_name_or_path=pretrained_name_or_path,
|
pretrained_name_or_path=pretrained_name_or_path,
|
||||||
config=config,
|
config=config,
|
||||||
@@ -85,6 +86,7 @@ class EVO1Policy(PreTrainedPolicy):
|
|||||||
"device": config.device,
|
"device": config.device,
|
||||||
"return_cls_only": config.return_cls_only,
|
"return_cls_only": config.return_cls_only,
|
||||||
"vlm_name": config.vlm_model_name,
|
"vlm_name": config.vlm_model_name,
|
||||||
|
"image_size": int(config.image_resolution[0]),
|
||||||
"vlm_num_layers": config.vlm_num_layers,
|
"vlm_num_layers": config.vlm_num_layers,
|
||||||
"vlm_dtype": config.vlm_dtype,
|
"vlm_dtype": config.vlm_dtype,
|
||||||
"use_flash_attn": config.use_flash_attn,
|
"use_flash_attn": config.use_flash_attn,
|
||||||
@@ -100,7 +102,8 @@ class EVO1Policy(PreTrainedPolicy):
|
|||||||
"dropout": config.dropout,
|
"dropout": config.dropout,
|
||||||
"num_inference_timesteps": config.num_inference_timesteps,
|
"num_inference_timesteps": config.num_inference_timesteps,
|
||||||
"num_categories": config.num_categories,
|
"num_categories": config.num_categories,
|
||||||
"enable_gradient_checkpointing": config.enable_gradient_checkpointing,
|
"enable_gradient_checkpointing": config.enable_gradient_checkpointing
|
||||||
|
and bool(config.finetune_vlm or config.finetune_language_model or config.finetune_vision_model),
|
||||||
"gradient_checkpointing_use_reentrant": config.gradient_checkpointing_use_reentrant,
|
"gradient_checkpointing_use_reentrant": config.gradient_checkpointing_use_reentrant,
|
||||||
"finetune_vlm": config.finetune_vlm,
|
"finetune_vlm": config.finetune_vlm,
|
||||||
"finetune_language_model": config.finetune_language_model,
|
"finetune_language_model": config.finetune_language_model,
|
||||||
@@ -303,6 +306,18 @@ class EVO1Policy(PreTrainedPolicy):
|
|||||||
or self.config.finetune_vision_model
|
or self.config.finetune_vision_model
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _keep_frozen_embedder_eval(self) -> None:
|
||||||
|
if self._tracks_vlm_gradients:
|
||||||
|
return
|
||||||
|
embedder = getattr(self.model, "embedder", None)
|
||||||
|
if embedder is not None:
|
||||||
|
embedder.eval()
|
||||||
|
|
||||||
|
def train(self, mode: bool = True):
|
||||||
|
super().train(mode)
|
||||||
|
self._keep_frozen_embedder_eval()
|
||||||
|
return self
|
||||||
|
|
||||||
def _collect_image_batches(self, batch: dict[str, Tensor]) -> tuple[list[list[Tensor]], Tensor]:
|
def _collect_image_batches(self, batch: dict[str, Tensor]) -> tuple[list[list[Tensor]], Tensor]:
|
||||||
camera_keys = self._camera_keys or sorted(key for key in batch if key.startswith(f"{OBS_IMAGES}."))
|
camera_keys = self._camera_keys or sorted(key for key in batch if key.startswith(f"{OBS_IMAGES}."))
|
||||||
if not camera_keys:
|
if not camera_keys:
|
||||||
@@ -348,23 +363,13 @@ class EVO1Policy(PreTrainedPolicy):
|
|||||||
) -> Tensor:
|
) -> Tensor:
|
||||||
track_vlm_gradients = self._tracks_vlm_gradients
|
track_vlm_gradients = self._tracks_vlm_gradients
|
||||||
grad_context = nullcontext() if track_vlm_gradients else torch.no_grad()
|
grad_context = nullcontext() if track_vlm_gradients else torch.no_grad()
|
||||||
embedder = getattr(self.model, "embedder", None)
|
with grad_context:
|
||||||
embedder_was_training = embedder.training if embedder is not None else None
|
fused_tokens = self.model.get_vl_embeddings(
|
||||||
|
images=image_batches,
|
||||||
if not track_vlm_gradients and embedder is not None:
|
image_mask=image_masks,
|
||||||
embedder.eval()
|
prompt=prompts,
|
||||||
|
return_cls_only=self.config.return_cls_only,
|
||||||
try:
|
)
|
||||||
with grad_context:
|
|
||||||
fused_tokens = self.model.get_vl_embeddings(
|
|
||||||
images=image_batches,
|
|
||||||
image_mask=image_masks,
|
|
||||||
prompt=prompts,
|
|
||||||
return_cls_only=self.config.return_cls_only,
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
if not track_vlm_gradients and embedder is not None and embedder_was_training is not None:
|
|
||||||
embedder.train(embedder_was_training)
|
|
||||||
|
|
||||||
if not track_vlm_gradients:
|
if not track_vlm_gradients:
|
||||||
fused_tokens = fused_tokens.detach()
|
fused_tokens = fused_tokens.detach()
|
||||||
@@ -439,7 +444,7 @@ class EVO1Policy(PreTrainedPolicy):
|
|||||||
embodiment_ids=embodiment_ids,
|
embodiment_ids=embodiment_ids,
|
||||||
)
|
)
|
||||||
actions = actions.view(states.shape[0], self.config.chunk_size, self.config.max_action_dim)
|
actions = actions.view(states.shape[0], self.config.chunk_size, self.config.max_action_dim)
|
||||||
return actions[:, :, : self._env_action_dim]
|
return actions
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
|
def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
|
||||||
|
|||||||
@@ -14,17 +14,24 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from copy import deepcopy
|
||||||
|
from dataclasses import dataclass
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||||
from lerobot.policies.evo1.configuration_evo1 import Evo1Config
|
from lerobot.policies.evo1.configuration_evo1 import Evo1Config
|
||||||
from lerobot.processor import (
|
from lerobot.processor import (
|
||||||
AddBatchDimensionProcessorStep,
|
AddBatchDimensionProcessorStep,
|
||||||
DeviceProcessorStep,
|
DeviceProcessorStep,
|
||||||
NormalizerProcessorStep,
|
NormalizerProcessorStep,
|
||||||
|
ObservationProcessorStep,
|
||||||
PolicyAction,
|
PolicyAction,
|
||||||
|
PolicyActionProcessorStep,
|
||||||
PolicyProcessorPipeline,
|
PolicyProcessorPipeline,
|
||||||
|
ProcessorStep,
|
||||||
|
ProcessorStepRegistry,
|
||||||
RenameObservationsProcessorStep,
|
RenameObservationsProcessorStep,
|
||||||
UnnormalizerProcessorStep,
|
UnnormalizerProcessorStep,
|
||||||
)
|
)
|
||||||
@@ -34,11 +41,13 @@ from lerobot.processor.converters import (
|
|||||||
policy_action_to_transition,
|
policy_action_to_transition,
|
||||||
transition_to_policy_action,
|
transition_to_policy_action,
|
||||||
)
|
)
|
||||||
|
from lerobot.types import EnvTransition, TransitionKey
|
||||||
from lerobot.utils.constants import (
|
from lerobot.utils.constants import (
|
||||||
ACTION,
|
ACTION,
|
||||||
DONE,
|
DONE,
|
||||||
INFO,
|
INFO,
|
||||||
OBS_PREFIX,
|
OBS_PREFIX,
|
||||||
|
OBS_STATE,
|
||||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||||
REWARD,
|
REWARD,
|
||||||
@@ -65,6 +74,305 @@ def evo1_batch_to_transition(batch: dict[str, Any]):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
@ProcessorStepRegistry.register(name="evo1_pad_state_processor")
|
||||||
|
class Evo1PadStateProcessorStep(ObservationProcessorStep):
|
||||||
|
"""Pad policy observations to EVO1's fixed state width before normalization."""
|
||||||
|
|
||||||
|
max_state_dim: int = 24
|
||||||
|
|
||||||
|
def observation(self, observation: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
if OBS_STATE not in observation:
|
||||||
|
return observation
|
||||||
|
|
||||||
|
state = observation[OBS_STATE]
|
||||||
|
state_dim = state.shape[-1]
|
||||||
|
if state_dim > self.max_state_dim:
|
||||||
|
raise ValueError(
|
||||||
|
f"EVO1 state has {state_dim} dims, which exceeds max_state_dim={self.max_state_dim}."
|
||||||
|
)
|
||||||
|
if state_dim < self.max_state_dim:
|
||||||
|
observation = observation.copy()
|
||||||
|
observation[OBS_STATE] = torch.nn.functional.pad(state, (0, self.max_state_dim - state_dim))
|
||||||
|
return observation
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
new_features = {ft: feats.copy() for ft, feats in features.items()}
|
||||||
|
state_feats = new_features.setdefault(FeatureType.STATE, {})
|
||||||
|
if OBS_STATE in state_feats:
|
||||||
|
state_feats[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.max_state_dim,))
|
||||||
|
return new_features
|
||||||
|
|
||||||
|
def get_config(self) -> dict[str, Any]:
|
||||||
|
return {"max_state_dim": self.max_state_dim}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
@ProcessorStepRegistry.register(name="evo1_pad_action_processor")
|
||||||
|
class Evo1PadActionProcessorStep(ProcessorStep):
|
||||||
|
"""Pad training actions and preserve the active action dimensions with action_mask."""
|
||||||
|
|
||||||
|
max_action_dim: int = 24
|
||||||
|
|
||||||
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||||
|
action = transition.get(TransitionKey.ACTION)
|
||||||
|
if action is None:
|
||||||
|
return transition
|
||||||
|
if not isinstance(action, PolicyAction):
|
||||||
|
raise ValueError(f"EVO1 action should be a PolicyAction tensor, but got {type(action)}.")
|
||||||
|
|
||||||
|
action_dim = action.shape[-1]
|
||||||
|
if action_dim > self.max_action_dim:
|
||||||
|
raise ValueError(
|
||||||
|
f"EVO1 action has {action_dim} dims, which exceeds max_action_dim={self.max_action_dim}."
|
||||||
|
)
|
||||||
|
|
||||||
|
new_transition = transition.copy()
|
||||||
|
new_action = action
|
||||||
|
if action_dim < self.max_action_dim:
|
||||||
|
new_action = torch.nn.functional.pad(action, (0, self.max_action_dim - action_dim))
|
||||||
|
|
||||||
|
complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
|
||||||
|
action_mask = complementary_data.get("action_mask")
|
||||||
|
if action_mask is None:
|
||||||
|
action_mask = torch.ones(action.shape, dtype=torch.bool, device=action.device)
|
||||||
|
else:
|
||||||
|
action_mask = torch.as_tensor(action_mask, dtype=torch.bool, device=action.device)
|
||||||
|
if action_mask.shape != action.shape:
|
||||||
|
raise ValueError(
|
||||||
|
f"action_mask shape {tuple(action_mask.shape)} does not match action shape {tuple(action.shape)}."
|
||||||
|
)
|
||||||
|
if action_dim < self.max_action_dim:
|
||||||
|
action_mask = torch.nn.functional.pad(action_mask, (0, self.max_action_dim - action_dim))
|
||||||
|
|
||||||
|
complementary_data["action_mask"] = action_mask
|
||||||
|
new_transition[TransitionKey.ACTION] = new_action
|
||||||
|
new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
|
||||||
|
return new_transition
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
new_features = {ft: feats.copy() for ft, feats in features.items()}
|
||||||
|
action_feats = new_features.setdefault(FeatureType.ACTION, {})
|
||||||
|
action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.max_action_dim,))
|
||||||
|
return new_features
|
||||||
|
|
||||||
|
def get_config(self) -> dict[str, Any]:
|
||||||
|
return {"max_action_dim": self.max_action_dim}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
@ProcessorStepRegistry.register(name="evo1_action_processor")
|
||||||
|
class Evo1ActionProcessorStep(PolicyActionProcessorStep):
|
||||||
|
"""Crop padded EVO1 actions and optionally binarize the LIBERO gripper channel."""
|
||||||
|
|
||||||
|
action_dim: int
|
||||||
|
binarize_gripper: bool = False
|
||||||
|
gripper_index: int = 6
|
||||||
|
gripper_threshold: float = 0.5
|
||||||
|
gripper_below_threshold_value: float = 1.0
|
||||||
|
gripper_above_threshold_value: float = -1.0
|
||||||
|
|
||||||
|
def action(self, action: PolicyAction) -> PolicyAction:
|
||||||
|
if action.shape[-1] < self.action_dim:
|
||||||
|
raise ValueError(
|
||||||
|
f"EVO1 action has {action.shape[-1]} dims, which is smaller than action_dim={self.action_dim}."
|
||||||
|
)
|
||||||
|
|
||||||
|
action = action[..., : self.action_dim]
|
||||||
|
if not self.binarize_gripper:
|
||||||
|
return action
|
||||||
|
|
||||||
|
if not 0 <= self.gripper_index < self.action_dim:
|
||||||
|
raise ValueError(
|
||||||
|
f"gripper_index={self.gripper_index} must be within action_dim={self.action_dim}."
|
||||||
|
)
|
||||||
|
|
||||||
|
action = action.clone()
|
||||||
|
below = torch.as_tensor(
|
||||||
|
self.gripper_below_threshold_value,
|
||||||
|
dtype=action.dtype,
|
||||||
|
device=action.device,
|
||||||
|
)
|
||||||
|
above = torch.as_tensor(
|
||||||
|
self.gripper_above_threshold_value,
|
||||||
|
dtype=action.dtype,
|
||||||
|
device=action.device,
|
||||||
|
)
|
||||||
|
action[..., self.gripper_index] = torch.where(
|
||||||
|
action[..., self.gripper_index] > self.gripper_threshold,
|
||||||
|
above,
|
||||||
|
below,
|
||||||
|
)
|
||||||
|
return action
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
new_features = {ft: feats.copy() for ft, feats in features.items()}
|
||||||
|
action_feats = new_features.setdefault(FeatureType.ACTION, {})
|
||||||
|
action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))
|
||||||
|
return new_features
|
||||||
|
|
||||||
|
def get_config(self) -> dict[str, Any]:
|
||||||
|
return {
|
||||||
|
"action_dim": self.action_dim,
|
||||||
|
"binarize_gripper": self.binarize_gripper,
|
||||||
|
"gripper_index": self.gripper_index,
|
||||||
|
"gripper_threshold": self.gripper_threshold,
|
||||||
|
"gripper_below_threshold_value": self.gripper_below_threshold_value,
|
||||||
|
"gripper_above_threshold_value": self.gripper_above_threshold_value,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _evo1_action_dim(config: Evo1Config) -> int:
|
||||||
|
if config.postprocess_action_dim is not None:
|
||||||
|
return config.postprocess_action_dim
|
||||||
|
action_feature = config.action_feature
|
||||||
|
if action_feature is None:
|
||||||
|
return config.max_action_dim
|
||||||
|
return int(action_feature.shape[0])
|
||||||
|
|
||||||
|
|
||||||
|
def _evo1_normalization_features(config: Evo1Config) -> dict[str, PolicyFeature]:
|
||||||
|
features = {**config.input_features, **config.output_features}
|
||||||
|
features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(config.max_state_dim,))
|
||||||
|
features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(config.max_action_dim,))
|
||||||
|
return features
|
||||||
|
|
||||||
|
|
||||||
|
def _evo1_action_features(config: Evo1Config) -> dict[str, PolicyFeature]:
|
||||||
|
return {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(config.max_action_dim,))}
|
||||||
|
|
||||||
|
|
||||||
|
_STAT_PAD_VALUES = {
|
||||||
|
"mean": 0.0,
|
||||||
|
"std": 1.0,
|
||||||
|
"min": -1.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"q01": -1.0,
|
||||||
|
"q99": 1.0,
|
||||||
|
"q10": -1.0,
|
||||||
|
"q90": 1.0,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _pad_stat_value(value: Any, target_dim: int, stat_name: str) -> torch.Tensor:
|
||||||
|
tensor = torch.as_tensor(value)
|
||||||
|
if not tensor.is_floating_point():
|
||||||
|
tensor = tensor.to(dtype=torch.float32)
|
||||||
|
if tensor.ndim == 0 or tensor.shape[-1] >= target_dim:
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
pad_shape = (*tensor.shape[:-1], target_dim - tensor.shape[-1])
|
||||||
|
pad_value = _STAT_PAD_VALUES.get(stat_name, 0.0)
|
||||||
|
padding = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device)
|
||||||
|
return torch.cat([tensor, padding], dim=-1)
|
||||||
|
|
||||||
|
|
||||||
|
def _pad_feature_stats(
|
||||||
|
stats: dict[str, dict[str, Any]],
|
||||||
|
feature_key: str,
|
||||||
|
target_dim: int,
|
||||||
|
) -> None:
|
||||||
|
if feature_key not in stats:
|
||||||
|
return
|
||||||
|
stats[feature_key] = {
|
||||||
|
stat_name: _pad_stat_value(stat_value, target_dim, stat_name)
|
||||||
|
for stat_name, stat_value in stats[feature_key].items()
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _pad_evo1_stats(
|
||||||
|
config: Evo1Config,
|
||||||
|
stats: dict[str, dict[str, Any]] | None,
|
||||||
|
) -> dict[str, dict[str, Any]] | None:
|
||||||
|
if stats is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
padded_stats = deepcopy(stats)
|
||||||
|
# Added dimensions represent zero-padding inside EVO1. These neutral stats keep
|
||||||
|
# padded observations at normalized zero and only provide shape compatibility.
|
||||||
|
_pad_feature_stats(padded_stats, OBS_STATE, config.max_state_dim)
|
||||||
|
_pad_feature_stats(padded_stats, ACTION, config.max_action_dim)
|
||||||
|
return padded_stats
|
||||||
|
|
||||||
|
|
||||||
|
def _refresh_evo1_normalization_steps(
|
||||||
|
config: Evo1Config,
|
||||||
|
preprocessor: PolicyProcessorPipeline,
|
||||||
|
postprocessor: PolicyProcessorPipeline,
|
||||||
|
) -> None:
|
||||||
|
normalization_features = _evo1_normalization_features(config)
|
||||||
|
action_features = _evo1_action_features(config)
|
||||||
|
|
||||||
|
for step in preprocessor.steps:
|
||||||
|
if isinstance(step, NormalizerProcessorStep):
|
||||||
|
step.features = normalization_features
|
||||||
|
step.stats = _pad_evo1_stats(config, step.stats)
|
||||||
|
step.to(device=step.device, dtype=step.dtype)
|
||||||
|
|
||||||
|
for step in postprocessor.steps:
|
||||||
|
if isinstance(step, UnnormalizerProcessorStep):
|
||||||
|
step.features = action_features
|
||||||
|
step.stats = _pad_evo1_stats(config, step.stats)
|
||||||
|
step.to(device=step.device, dtype=step.dtype)
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_evo1_processor_steps(
|
||||||
|
config: Evo1Config,
|
||||||
|
preprocessor: PolicyProcessorPipeline,
|
||||||
|
postprocessor: PolicyProcessorPipeline,
|
||||||
|
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
||||||
|
"""Add EVO1 processor steps when loading older checkpoints that do not serialize them."""
|
||||||
|
|
||||||
|
has_state_padding = any(isinstance(step, Evo1PadStateProcessorStep) for step in preprocessor.steps)
|
||||||
|
if not has_state_padding:
|
||||||
|
steps = list(preprocessor.steps)
|
||||||
|
insert_idx = next(
|
||||||
|
(idx for idx, step in enumerate(steps) if isinstance(step, NormalizerProcessorStep)),
|
||||||
|
len(steps),
|
||||||
|
)
|
||||||
|
steps.insert(insert_idx, Evo1PadStateProcessorStep(max_state_dim=config.max_state_dim))
|
||||||
|
preprocessor.steps = steps
|
||||||
|
|
||||||
|
has_action_padding = any(isinstance(step, Evo1PadActionProcessorStep) for step in preprocessor.steps)
|
||||||
|
if not has_action_padding:
|
||||||
|
steps = list(preprocessor.steps)
|
||||||
|
insert_idx = next(
|
||||||
|
(idx for idx, step in enumerate(steps) if isinstance(step, NormalizerProcessorStep)),
|
||||||
|
len(steps),
|
||||||
|
)
|
||||||
|
steps.insert(insert_idx, Evo1PadActionProcessorStep(max_action_dim=config.max_action_dim))
|
||||||
|
preprocessor.steps = steps
|
||||||
|
|
||||||
|
has_action_processor = any(isinstance(step, Evo1ActionProcessorStep) for step in postprocessor.steps)
|
||||||
|
if not has_action_processor:
|
||||||
|
steps = list(postprocessor.steps)
|
||||||
|
insert_idx = next(
|
||||||
|
(idx + 1 for idx, step in enumerate(steps) if isinstance(step, UnnormalizerProcessorStep)),
|
||||||
|
0,
|
||||||
|
)
|
||||||
|
steps.insert(
|
||||||
|
insert_idx,
|
||||||
|
Evo1ActionProcessorStep(
|
||||||
|
action_dim=_evo1_action_dim(config),
|
||||||
|
binarize_gripper=config.binarize_gripper,
|
||||||
|
gripper_index=config.gripper_index,
|
||||||
|
gripper_threshold=config.gripper_threshold,
|
||||||
|
gripper_below_threshold_value=config.gripper_below_threshold_value,
|
||||||
|
gripper_above_threshold_value=config.gripper_above_threshold_value,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
postprocessor.steps = steps
|
||||||
|
|
||||||
|
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
|
||||||
|
return preprocessor, postprocessor
|
||||||
|
|
||||||
|
|
||||||
def make_evo1_pre_post_processors(
|
def make_evo1_pre_post_processors(
|
||||||
config: Evo1Config,
|
config: Evo1Config,
|
||||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||||
@@ -72,21 +380,35 @@ def make_evo1_pre_post_processors(
|
|||||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||||
]:
|
]:
|
||||||
|
normalization_features = _evo1_normalization_features(config)
|
||||||
|
action_features = _evo1_action_features(config)
|
||||||
|
normalization_stats = _pad_evo1_stats(config, dataset_stats)
|
||||||
|
|
||||||
input_steps = [
|
input_steps = [
|
||||||
RenameObservationsProcessorStep(rename_map={}),
|
RenameObservationsProcessorStep(rename_map={}),
|
||||||
AddBatchDimensionProcessorStep(),
|
AddBatchDimensionProcessorStep(),
|
||||||
|
Evo1PadStateProcessorStep(max_state_dim=config.max_state_dim),
|
||||||
|
Evo1PadActionProcessorStep(max_action_dim=config.max_action_dim),
|
||||||
NormalizerProcessorStep(
|
NormalizerProcessorStep(
|
||||||
features={**config.input_features, **config.output_features},
|
features=normalization_features,
|
||||||
norm_map=config.normalization_mapping,
|
norm_map=config.normalization_mapping,
|
||||||
stats=dataset_stats,
|
stats=normalization_stats,
|
||||||
),
|
),
|
||||||
DeviceProcessorStep(device=config.device),
|
DeviceProcessorStep(device=config.device),
|
||||||
]
|
]
|
||||||
output_steps = [
|
output_steps = [
|
||||||
UnnormalizerProcessorStep(
|
UnnormalizerProcessorStep(
|
||||||
features=config.output_features,
|
features=action_features,
|
||||||
norm_map=config.normalization_mapping,
|
norm_map=config.normalization_mapping,
|
||||||
stats=dataset_stats,
|
stats=normalization_stats,
|
||||||
|
),
|
||||||
|
Evo1ActionProcessorStep(
|
||||||
|
action_dim=_evo1_action_dim(config),
|
||||||
|
binarize_gripper=config.binarize_gripper,
|
||||||
|
gripper_index=config.gripper_index,
|
||||||
|
gripper_threshold=config.gripper_threshold,
|
||||||
|
gripper_below_threshold_value=config.gripper_below_threshold_value,
|
||||||
|
gripper_above_threshold_value=config.gripper_above_threshold_value,
|
||||||
),
|
),
|
||||||
DeviceProcessorStep(device="cpu"),
|
DeviceProcessorStep(device="cpu"),
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -311,6 +311,14 @@ def make_pre_post_processors(
|
|||||||
to_output=transition_to_policy_action,
|
to_output=transition_to_policy_action,
|
||||||
)
|
)
|
||||||
_reconnect_relative_absolute_steps(preprocessor, postprocessor)
|
_reconnect_relative_absolute_steps(preprocessor, postprocessor)
|
||||||
|
if isinstance(policy_cfg, Evo1Config):
|
||||||
|
from .evo1.processor_evo1 import ensure_evo1_processor_steps
|
||||||
|
|
||||||
|
preprocessor, postprocessor = ensure_evo1_processor_steps(
|
||||||
|
policy_cfg,
|
||||||
|
preprocessor,
|
||||||
|
postprocessor,
|
||||||
|
)
|
||||||
return preprocessor, postprocessor
|
return preprocessor, postprocessor
|
||||||
|
|
||||||
# Create a new processor based on policy type
|
# Create a new processor based on policy type
|
||||||
|
|||||||
@@ -40,7 +40,7 @@ from .converters import (
|
|||||||
)
|
)
|
||||||
from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
|
from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
|
||||||
from .device_processor import DeviceProcessorStep
|
from .device_processor import DeviceProcessorStep
|
||||||
from .env_processor import IsaaclabArenaProcessorStep, LiberoActionProcessorStep, LiberoProcessorStep
|
from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
|
||||||
from .factory import (
|
from .factory import (
|
||||||
make_default_processors,
|
make_default_processors,
|
||||||
make_default_robot_action_processor,
|
make_default_robot_action_processor,
|
||||||
@@ -149,7 +149,6 @@ __all__ = [
|
|||||||
"RewardProcessorStep",
|
"RewardProcessorStep",
|
||||||
"DataProcessorPipeline",
|
"DataProcessorPipeline",
|
||||||
"IsaaclabArenaProcessorStep",
|
"IsaaclabArenaProcessorStep",
|
||||||
"LiberoActionProcessorStep",
|
|
||||||
"LiberoProcessorStep",
|
"LiberoProcessorStep",
|
||||||
"TimeLimitProcessorStep",
|
"TimeLimitProcessorStep",
|
||||||
"AddBatchDimensionProcessorStep",
|
"AddBatchDimensionProcessorStep",
|
||||||
|
|||||||
@@ -18,9 +18,9 @@ from dataclasses import dataclass
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||||
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_PREFIX, OBS_STATE, OBS_STR
|
from lerobot.utils.constants import OBS_IMAGES, OBS_PREFIX, OBS_STATE, OBS_STR
|
||||||
|
|
||||||
from .pipeline import ActionProcessorStep, ObservationProcessorStep, ProcessorStepRegistry
|
from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -46,8 +46,6 @@ class LiberoProcessorStep(ObservationProcessorStep):
|
|||||||
- This accounts for the HuggingFaceVLA/libero camera orientation convention.
|
- This accounts for the HuggingFaceVLA/libero camera orientation convention.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
max_state_dim: int | None = None
|
|
||||||
|
|
||||||
def _process_observation(self, observation):
|
def _process_observation(self, observation):
|
||||||
"""
|
"""
|
||||||
Processes both image and robot_state observations from LIBERO.
|
Processes both image and robot_state observations from LIBERO.
|
||||||
@@ -80,16 +78,6 @@ class LiberoProcessorStep(ObservationProcessorStep):
|
|||||||
state = state.float()
|
state = state.float()
|
||||||
if state.dim() == 1:
|
if state.dim() == 1:
|
||||||
state = state.unsqueeze(0)
|
state = state.unsqueeze(0)
|
||||||
if self.max_state_dim is not None:
|
|
||||||
if state.shape[-1] > self.max_state_dim:
|
|
||||||
raise ValueError(
|
|
||||||
f"LIBERO state has {state.shape[-1]} dims, which is larger than "
|
|
||||||
f"configured max_state_dim={self.max_state_dim}."
|
|
||||||
)
|
|
||||||
if state.shape[-1] < self.max_state_dim:
|
|
||||||
pad_width = self.max_state_dim - state.shape[-1]
|
|
||||||
state = torch.nn.functional.pad(state, (0, pad_width))
|
|
||||||
|
|
||||||
processed_obs[OBS_STATE] = state
|
processed_obs[OBS_STATE] = state
|
||||||
return processed_obs
|
return processed_obs
|
||||||
|
|
||||||
@@ -112,7 +100,7 @@ class LiberoProcessorStep(ObservationProcessorStep):
|
|||||||
# add our new flattened state
|
# add our new flattened state
|
||||||
state_feats[OBS_STATE] = PolicyFeature(
|
state_feats[OBS_STATE] = PolicyFeature(
|
||||||
type=FeatureType.STATE,
|
type=FeatureType.STATE,
|
||||||
shape=(self.max_state_dim or 8,), # [eef_pos(3), axis_angle(3), gripper(2)] plus padding
|
shape=(8,), # [eef_pos(3), axis_angle(3), gripper(2)]
|
||||||
)
|
)
|
||||||
|
|
||||||
new_features[FeatureType.STATE] = state_feats
|
new_features[FeatureType.STATE] = state_feats
|
||||||
@@ -122,9 +110,6 @@ class LiberoProcessorStep(ObservationProcessorStep):
|
|||||||
def observation(self, observation):
|
def observation(self, observation):
|
||||||
return self._process_observation(observation)
|
return self._process_observation(observation)
|
||||||
|
|
||||||
def get_config(self) -> dict:
|
|
||||||
return {"max_state_dim": self.max_state_dim}
|
|
||||||
|
|
||||||
def _quat2axisangle(self, quat: torch.Tensor) -> torch.Tensor:
|
def _quat2axisangle(self, quat: torch.Tensor) -> torch.Tensor:
|
||||||
"""
|
"""
|
||||||
Convert batched quaternions to axis-angle format.
|
Convert batched quaternions to axis-angle format.
|
||||||
@@ -167,68 +152,6 @@ class LiberoProcessorStep(ObservationProcessorStep):
|
|||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
@ProcessorStepRegistry.register(name="libero_action_processor")
|
|
||||||
class LiberoActionProcessorStep(ActionProcessorStep):
|
|
||||||
"""Slices padded policy actions back to the executable LIBERO action space."""
|
|
||||||
|
|
||||||
action_dim: int = 7
|
|
||||||
binarize_gripper: bool = False
|
|
||||||
gripper_index: int = 6
|
|
||||||
gripper_threshold: float = 0.5
|
|
||||||
gripper_below_threshold_value: float = 1.0
|
|
||||||
gripper_above_threshold_value: float = -1.0
|
|
||||||
|
|
||||||
def action(self, action):
|
|
||||||
if action.shape[-1] < self.action_dim:
|
|
||||||
raise ValueError(
|
|
||||||
f"LIBERO action has {action.shape[-1]} dims, which is smaller than action_dim={self.action_dim}."
|
|
||||||
)
|
|
||||||
action = action[..., : self.action_dim]
|
|
||||||
if not self.binarize_gripper:
|
|
||||||
return action
|
|
||||||
|
|
||||||
if not 0 <= self.gripper_index < self.action_dim:
|
|
||||||
raise ValueError(
|
|
||||||
f"gripper_index={self.gripper_index} must be within sliced action_dim={self.action_dim}."
|
|
||||||
)
|
|
||||||
action = action.clone()
|
|
||||||
below = torch.as_tensor(
|
|
||||||
self.gripper_below_threshold_value,
|
|
||||||
dtype=action.dtype,
|
|
||||||
device=action.device,
|
|
||||||
)
|
|
||||||
above = torch.as_tensor(
|
|
||||||
self.gripper_above_threshold_value,
|
|
||||||
dtype=action.dtype,
|
|
||||||
device=action.device,
|
|
||||||
)
|
|
||||||
action[..., self.gripper_index] = torch.where(
|
|
||||||
action[..., self.gripper_index] > self.gripper_threshold,
|
|
||||||
above,
|
|
||||||
below,
|
|
||||||
)
|
|
||||||
return action
|
|
||||||
|
|
||||||
def transform_features(
|
|
||||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
|
||||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
|
||||||
new_features = {ft: feats.copy() for ft, feats in features.items()}
|
|
||||||
action_feats = new_features.setdefault(FeatureType.ACTION, {})
|
|
||||||
action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))
|
|
||||||
return new_features
|
|
||||||
|
|
||||||
def get_config(self) -> dict:
|
|
||||||
return {
|
|
||||||
"action_dim": self.action_dim,
|
|
||||||
"binarize_gripper": self.binarize_gripper,
|
|
||||||
"gripper_index": self.gripper_index,
|
|
||||||
"gripper_threshold": self.gripper_threshold,
|
|
||||||
"gripper_below_threshold_value": self.gripper_below_threshold_value,
|
|
||||||
"gripper_above_threshold_value": self.gripper_above_threshold_value,
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@ProcessorStepRegistry.register(name="isaaclab_arena_processor")
|
@ProcessorStepRegistry.register(name="isaaclab_arena_processor")
|
||||||
class IsaaclabArenaProcessorStep(ObservationProcessorStep):
|
class IsaaclabArenaProcessorStep(ObservationProcessorStep):
|
||||||
|
|||||||
+10
-61
@@ -13,7 +13,7 @@ from gymnasium.envs.registration import register, registry as gym_registry
|
|||||||
from lerobot.configs.types import PolicyFeature
|
from lerobot.configs.types import PolicyFeature
|
||||||
from lerobot.envs.configs import EnvConfig, LiberoEnv
|
from lerobot.envs.configs import EnvConfig, LiberoEnv
|
||||||
from lerobot.envs.factory import make_env, make_env_config, make_env_pre_post_processors
|
from lerobot.envs.factory import make_env, make_env_config, make_env_pre_post_processors
|
||||||
from lerobot.processor import LiberoActionProcessorStep, LiberoProcessorStep
|
from lerobot.processor import LiberoProcessorStep
|
||||||
from lerobot.utils.constants import OBS_PREFIX, OBS_STATE
|
from lerobot.utils.constants import OBS_PREFIX, OBS_STATE
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -86,38 +86,18 @@ def test_processors_delegation_supports_legacy_override_signature():
|
|||||||
assert isinstance(post, DataProcessorPipeline)
|
assert isinstance(post, DataProcessorPipeline)
|
||||||
|
|
||||||
|
|
||||||
def test_libero_evo1_processors_use_padded_state_and_env_action_dim():
|
def test_libero_processors_are_policy_agnostic():
|
||||||
"""EVO1 uses padded LIBERO state features while env actions stay executable."""
|
|
||||||
|
|
||||||
class _Evo1Config:
|
|
||||||
type = "evo1"
|
|
||||||
max_state_dim = 24
|
|
||||||
|
|
||||||
cfg = LiberoEnv()
|
cfg = LiberoEnv()
|
||||||
pre, post = make_env_pre_post_processors(cfg, policy_cfg=_Evo1Config())
|
pre, post = make_env_pre_post_processors(cfg, policy_cfg=object())
|
||||||
|
|
||||||
assert isinstance(pre.steps[0], LiberoProcessorStep)
|
assert isinstance(pre.steps[0], LiberoProcessorStep)
|
||||||
assert pre.steps[0].max_state_dim == 24
|
assert len(post.steps) == 0
|
||||||
assert isinstance(post.steps[0], LiberoActionProcessorStep)
|
|
||||||
assert post.steps[0].action_dim == cfg.features["action"].shape[0] == 7
|
|
||||||
assert post.steps[0].binarize_gripper is True
|
|
||||||
|
|
||||||
class _OtherConfig:
|
|
||||||
type = "other"
|
|
||||||
|
|
||||||
pre_other, post_other = make_env_pre_post_processors(cfg, policy_cfg=_OtherConfig())
|
|
||||||
assert pre_other.steps[0].max_state_dim is None
|
|
||||||
assert post_other.steps[0].binarize_gripper is False
|
|
||||||
|
|
||||||
cfg.binarize_gripper = False
|
|
||||||
_, post_disabled = make_env_pre_post_processors(cfg, policy_cfg=_Evo1Config())
|
|
||||||
assert post_disabled.steps[0].binarize_gripper is False
|
|
||||||
|
|
||||||
|
|
||||||
def test_libero_processor_pads_state_to_max_dim():
|
def test_libero_processor_flattens_state_to_raw_8_dim():
|
||||||
step = LiberoProcessorStep(max_state_dim=24)
|
step = LiberoProcessorStep()
|
||||||
observation = {
|
observation = {
|
||||||
OBS_PREFIX
|
OBS_PREFIX + "robot_state": {
|
||||||
+ "robot_state": {
|
|
||||||
"eef": {
|
"eef": {
|
||||||
"pos": torch.tensor([[1.0, 2.0, 3.0]]),
|
"pos": torch.tensor([[1.0, 2.0, 3.0]]),
|
||||||
"quat": torch.tensor([[0.0, 0.0, 0.0, 1.0]]),
|
"quat": torch.tensor([[0.0, 0.0, 0.0, 1.0]]),
|
||||||
@@ -127,39 +107,8 @@ def test_libero_processor_pads_state_to_max_dim():
|
|||||||
}
|
}
|
||||||
|
|
||||||
state = step.observation(observation)[OBS_STATE]
|
state = step.observation(observation)[OBS_STATE]
|
||||||
assert state.shape == (1, 24)
|
assert state.shape == (1, 8)
|
||||||
assert torch.allclose(state[:, :8], torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 4.0, 5.0]]))
|
assert torch.allclose(state, torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 4.0, 5.0]]))
|
||||||
assert torch.count_nonzero(state[:, 8:]).item() == 0
|
|
||||||
|
|
||||||
|
|
||||||
def test_libero_action_processor_slices_padded_action():
|
|
||||||
step = LiberoActionProcessorStep(action_dim=7)
|
|
||||||
action = torch.arange(2 * 3 * 24, dtype=torch.float32).reshape(2, 3, 24)
|
|
||||||
|
|
||||||
sliced = step.action(action)
|
|
||||||
assert sliced.shape == (2, 3, 7)
|
|
||||||
assert torch.equal(sliced, action[..., :7])
|
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="smaller than action_dim=7"):
|
|
||||||
step.action(torch.zeros(2, 6))
|
|
||||||
|
|
||||||
|
|
||||||
def test_libero_action_processor_can_binarize_gripper():
|
|
||||||
step = LiberoActionProcessorStep(action_dim=7, binarize_gripper=True)
|
|
||||||
action = torch.tensor(
|
|
||||||
[
|
|
||||||
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.5, 7.0],
|
|
||||||
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.6, 7.0],
|
|
||||||
],
|
|
||||||
dtype=torch.float32,
|
|
||||||
)
|
|
||||||
|
|
||||||
processed = step.action(action)
|
|
||||||
|
|
||||||
assert processed.shape == (2, 7)
|
|
||||||
assert torch.equal(processed[:, :6], action[:, :6])
|
|
||||||
assert torch.equal(processed[:, 6], torch.tensor([1.0, -1.0]))
|
|
||||||
assert torch.equal(action[:, 6], torch.tensor([0.5, 0.6]))
|
|
||||||
|
|
||||||
|
|
||||||
def test_base_create_envs():
|
def test_base_create_envs():
|
||||||
|
|||||||
@@ -16,6 +16,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
import torch
|
import torch
|
||||||
from torch import nn
|
from torch import nn
|
||||||
|
|
||||||
@@ -23,7 +24,15 @@ import lerobot.policies.evo1.modeling_evo1 as modeling_evo1
|
|||||||
from lerobot.configs.types import FeatureType, PolicyFeature
|
from lerobot.configs.types import FeatureType, PolicyFeature
|
||||||
from lerobot.policies.evo1.configuration_evo1 import Evo1Config
|
from lerobot.policies.evo1.configuration_evo1 import Evo1Config
|
||||||
from lerobot.policies.evo1.flow_matching import FlowmatchingActionHead
|
from lerobot.policies.evo1.flow_matching import FlowmatchingActionHead
|
||||||
|
from lerobot.policies.evo1.processor_evo1 import (
|
||||||
|
Evo1ActionProcessorStep,
|
||||||
|
Evo1PadActionProcessorStep,
|
||||||
|
Evo1PadStateProcessorStep,
|
||||||
|
ensure_evo1_processor_steps,
|
||||||
|
make_evo1_pre_post_processors,
|
||||||
|
)
|
||||||
from lerobot.policies.factory import get_policy_class, make_policy_config
|
from lerobot.policies.factory import get_policy_class, make_policy_config
|
||||||
|
from lerobot.processor import NormalizerProcessorStep, PolicyProcessorPipeline, UnnormalizerProcessorStep
|
||||||
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
|
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
|
||||||
|
|
||||||
STATE_DIM = 4
|
STATE_DIM = 4
|
||||||
@@ -108,6 +117,19 @@ def make_batch(include_action=True):
|
|||||||
return batch
|
return batch
|
||||||
|
|
||||||
|
|
||||||
|
def make_stats(state_dim=STATE_DIM, action_dim=ACTION_DIM):
|
||||||
|
return {
|
||||||
|
OBS_STATE: {
|
||||||
|
"min": torch.full((state_dim,), -2.0),
|
||||||
|
"max": torch.full((state_dim,), 2.0),
|
||||||
|
},
|
||||||
|
ACTION: {
|
||||||
|
"min": torch.full((action_dim,), -1.0),
|
||||||
|
"max": torch.full((action_dim,), 1.0),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def test_evo1_factory_registration():
|
def test_evo1_factory_registration():
|
||||||
cfg = make_policy_config(
|
cfg = make_policy_config(
|
||||||
"evo1",
|
"evo1",
|
||||||
@@ -191,22 +213,151 @@ def test_evo1_stage_defaults_and_consistency():
|
|||||||
raise AssertionError("Expected inconsistent finetune config to raise ValueError")
|
raise AssertionError("Expected inconsistent finetune config to raise ValueError")
|
||||||
|
|
||||||
|
|
||||||
|
def test_evo1_rejects_non_square_image_resolution():
|
||||||
|
with pytest.raises(ValueError, match="square image_resolution"):
|
||||||
|
make_config(image_resolution=(448, 320))
|
||||||
|
|
||||||
|
|
||||||
|
def test_evo1_build_model_config_uses_image_resolution_and_trainable_checkpointing():
|
||||||
|
stage1 = make_config(training_stage="stage1", image_resolution=(224, 224))
|
||||||
|
stage1_model_config = modeling_evo1.EVO1Policy._build_model_config(stage1)
|
||||||
|
|
||||||
|
assert stage1_model_config["image_size"] == 224
|
||||||
|
assert stage1_model_config["enable_gradient_checkpointing"] is False
|
||||||
|
|
||||||
|
stage2 = make_config(training_stage="stage2", image_resolution=(224, 224))
|
||||||
|
stage2_model_config = modeling_evo1.EVO1Policy._build_model_config(stage2)
|
||||||
|
|
||||||
|
assert stage2_model_config["enable_gradient_checkpointing"] is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_evo1_policy_processors_pad_state_crop_action_and_binarize_gripper():
|
||||||
|
libero_action_dim = 7
|
||||||
|
config = make_config(
|
||||||
|
max_state_dim=MAX_STATE_DIM,
|
||||||
|
max_action_dim=8,
|
||||||
|
postprocess_action_dim=libero_action_dim,
|
||||||
|
binarize_gripper=True,
|
||||||
|
output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(libero_action_dim,))},
|
||||||
|
)
|
||||||
|
stats = make_stats(action_dim=libero_action_dim)
|
||||||
|
|
||||||
|
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=stats)
|
||||||
|
|
||||||
|
assert isinstance(preprocessor.steps[2], Evo1PadStateProcessorStep)
|
||||||
|
assert isinstance(preprocessor.steps[3], Evo1PadActionProcessorStep)
|
||||||
|
assert isinstance(preprocessor.steps[4], NormalizerProcessorStep)
|
||||||
|
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||||
|
assert isinstance(postprocessor.steps[1], Evo1ActionProcessorStep)
|
||||||
|
|
||||||
|
normalizer = preprocessor.steps[4]
|
||||||
|
assert normalizer.features[OBS_STATE].shape == (MAX_STATE_DIM,)
|
||||||
|
assert normalizer.features[ACTION].shape == (8,)
|
||||||
|
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
|
||||||
|
assert normalizer._tensor_stats[ACTION]["min"].shape == (8,)
|
||||||
|
|
||||||
|
processed_batch = preprocessor(
|
||||||
|
{
|
||||||
|
"task": "pick the block",
|
||||||
|
OBS_STATE: torch.zeros(STATE_DIM),
|
||||||
|
ACTION: torch.zeros(libero_action_dim),
|
||||||
|
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
processed_state = processed_batch[OBS_STATE]
|
||||||
|
assert processed_state.shape == (1, MAX_STATE_DIM)
|
||||||
|
assert torch.allclose(processed_state, torch.zeros_like(processed_state))
|
||||||
|
assert processed_batch[ACTION].shape == (1, 8)
|
||||||
|
assert torch.allclose(processed_batch[ACTION], torch.zeros_like(processed_batch[ACTION]))
|
||||||
|
assert processed_batch["action_mask"].shape == (1, 8)
|
||||||
|
assert processed_batch["action_mask"][:, :libero_action_dim].all()
|
||||||
|
assert not processed_batch["action_mask"][:, libero_action_dim:].any()
|
||||||
|
|
||||||
|
action = torch.tensor(
|
||||||
|
[
|
||||||
|
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.5, 0.7],
|
||||||
|
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7],
|
||||||
|
],
|
||||||
|
dtype=torch.float32,
|
||||||
|
)
|
||||||
|
processed = postprocessor(action)
|
||||||
|
|
||||||
|
assert processed.shape == (2, 7)
|
||||||
|
assert torch.allclose(processed[:, :6], action[:, :6])
|
||||||
|
assert torch.equal(processed[:, 6], torch.tensor([1.0, -1.0]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_evo1_legacy_processors_are_completed_before_normalization():
|
||||||
|
config = make_config(
|
||||||
|
max_state_dim=MAX_STATE_DIM,
|
||||||
|
max_action_dim=8,
|
||||||
|
postprocess_action_dim=7,
|
||||||
|
binarize_gripper=True,
|
||||||
|
)
|
||||||
|
stats = make_stats(action_dim=7)
|
||||||
|
legacy_pre = PolicyProcessorPipeline(
|
||||||
|
steps=[
|
||||||
|
NormalizerProcessorStep(
|
||||||
|
features={**config.input_features, **config.output_features},
|
||||||
|
norm_map=config.normalization_mapping,
|
||||||
|
stats=stats,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
legacy_post = PolicyProcessorPipeline(
|
||||||
|
steps=[
|
||||||
|
UnnormalizerProcessorStep(
|
||||||
|
features=config.output_features,
|
||||||
|
norm_map=config.normalization_mapping,
|
||||||
|
stats=stats,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
preprocessor, postprocessor = ensure_evo1_processor_steps(config, legacy_pre, legacy_post)
|
||||||
|
|
||||||
|
assert isinstance(preprocessor.steps[0], Evo1PadStateProcessorStep)
|
||||||
|
assert isinstance(preprocessor.steps[1], Evo1PadActionProcessorStep)
|
||||||
|
assert isinstance(preprocessor.steps[2], NormalizerProcessorStep)
|
||||||
|
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||||
|
assert isinstance(postprocessor.steps[1], Evo1ActionProcessorStep)
|
||||||
|
assert postprocessor.steps[1].action_dim == 7
|
||||||
|
assert postprocessor.steps[1].binarize_gripper is True
|
||||||
|
assert preprocessor.steps[2].features[OBS_STATE].shape == (MAX_STATE_DIM,)
|
||||||
|
assert preprocessor.steps[2]._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
|
||||||
|
assert preprocessor.steps[2]._tensor_stats[ACTION]["min"].shape == (8,)
|
||||||
|
assert postprocessor.steps[0].features[ACTION].shape == (8,)
|
||||||
|
assert postprocessor.steps[0]._tensor_stats[ACTION]["min"].shape == (8,)
|
||||||
|
|
||||||
|
preprocessor, postprocessor = ensure_evo1_processor_steps(config, preprocessor, postprocessor)
|
||||||
|
assert sum(isinstance(step, Evo1PadStateProcessorStep) for step in preprocessor.steps) == 1
|
||||||
|
assert sum(isinstance(step, Evo1PadActionProcessorStep) for step in preprocessor.steps) == 1
|
||||||
|
assert sum(isinstance(step, Evo1ActionProcessorStep) for step in postprocessor.steps) == 1
|
||||||
|
|
||||||
|
|
||||||
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
|
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
|
||||||
monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
|
monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
|
||||||
policy = modeling_evo1.EVO1Policy(make_config())
|
policy = modeling_evo1.EVO1Policy(make_config())
|
||||||
|
preprocessor, _postprocessor = make_evo1_pre_post_processors(policy.config, dataset_stats=make_stats())
|
||||||
|
training_batch = preprocessor(make_batch(include_action=True))
|
||||||
|
|
||||||
loss, metrics = policy.forward(make_batch(include_action=True))
|
assert training_batch[ACTION].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
|
||||||
|
assert training_batch["action_mask"].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
|
||||||
|
assert training_batch["action_mask"][:, :, :ACTION_DIM].all()
|
||||||
|
assert not training_batch["action_mask"][:, :, ACTION_DIM:].any()
|
||||||
|
|
||||||
|
loss, metrics = policy.forward(training_batch)
|
||||||
assert loss.ndim == 0
|
assert loss.ndim == 0
|
||||||
assert torch.isfinite(loss)
|
assert torch.isfinite(loss)
|
||||||
assert metrics["active_action_dims"] == ACTION_DIM * CHUNK_SIZE
|
assert metrics["active_action_dims"] == ACTION_DIM * CHUNK_SIZE
|
||||||
assert policy.model.get_vl_embeddings_calls == 1
|
assert policy.model.get_vl_embeddings_calls == 1
|
||||||
|
|
||||||
action_chunk = policy.predict_action_chunk(make_batch(include_action=False))
|
action_chunk = policy.predict_action_chunk(make_batch(include_action=False))
|
||||||
assert action_chunk.shape == (2, CHUNK_SIZE, ACTION_DIM)
|
assert action_chunk.shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
|
||||||
|
|
||||||
policy.reset()
|
policy.reset()
|
||||||
selected = policy.select_action(make_batch(include_action=False))
|
selected = policy.select_action(make_batch(include_action=False))
|
||||||
assert selected.shape == (2, ACTION_DIM)
|
assert selected.shape == (2, MAX_ACTION_DIM)
|
||||||
|
|
||||||
|
|
||||||
def test_stage1_frozen_vlm_embeddings_do_not_track_gradients(monkeypatch):
|
def test_stage1_frozen_vlm_embeddings_do_not_track_gradients(monkeypatch):
|
||||||
@@ -220,7 +371,7 @@ def test_stage1_frozen_vlm_embeddings_do_not_track_gradients(monkeypatch):
|
|||||||
assert policy.model.grad_enabled_calls == [False]
|
assert policy.model.grad_enabled_calls == [False]
|
||||||
assert policy.model.embedder_training_calls == [False]
|
assert policy.model.embedder_training_calls == [False]
|
||||||
assert not fused_tokens.requires_grad
|
assert not fused_tokens.requires_grad
|
||||||
assert policy.model.embedder.training is True
|
assert policy.model.embedder.training is False
|
||||||
|
|
||||||
|
|
||||||
def test_stage2_vlm_embeddings_track_gradients(monkeypatch):
|
def test_stage2_vlm_embeddings_track_gradients(monkeypatch):
|
||||||
|
|||||||
Reference in New Issue
Block a user