Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,8 @@
|
||||
"""Environment registrations for harrix."""
|
||||
|
||||
_LIBERO_IMPORT_ERROR = None
|
||||
|
||||
try:
|
||||
from wall_x._vendor.harrix.envs import libero # noqa: F401
|
||||
except ModuleNotFoundError as exc:
|
||||
_LIBERO_IMPORT_ERROR = exc
|
||||
@@ -0,0 +1,119 @@
|
||||
"""Base environment abstraction.
|
||||
|
||||
Two execution granularities are supported:
|
||||
|
||||
- Episode-level ``run_episode``: caller supplies a predict callback and the env
|
||||
owns the full episode loop.
|
||||
- Chunk-level ``reset_episode`` + ``execute_chunk``: caller runs the model
|
||||
between chunks and feeds action chunks back to the env.
|
||||
|
||||
Subclasses must implement the chunk-level primitives. The default episode loop
|
||||
is built on top of those primitives.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Callable
|
||||
|
||||
import numpy as np
|
||||
|
||||
from wall_x._vendor.harrix.eval_config import EvalConfig
|
||||
|
||||
|
||||
class BaseEnv(ABC):
|
||||
|
||||
@abstractmethod
|
||||
def __init__(self, cfg: EvalConfig, worker_id: int) -> None:
|
||||
"""Perform env-specific setup."""
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def enumerate_episodes(cls, cfg: EvalConfig) -> list[tuple]:
|
||||
"""Return episode ids to seed JobState before env instances are built.
|
||||
|
||||
JobState treats the returned tuples as opaque ids.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def robot_spec(self) -> dict:
|
||||
"""Return robot metadata for driver/checkpoint validation.
|
||||
|
||||
Fields:
|
||||
- dof_layout: dict[str, int]
|
||||
- cam_names: list[str]
|
||||
- norm_key: str
|
||||
"""
|
||||
|
||||
# ---- chunk-level primitives ----
|
||||
|
||||
@abstractmethod
|
||||
def reset_episode(self, ep_id: tuple) -> dict:
|
||||
"""Start a new episode and return the first fresh observation.
|
||||
|
||||
Returns {"obs": dict, "instruction": str, "task_desc": str (optional)}.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def execute_chunk(self, actions: np.ndarray) -> dict:
|
||||
"""Execute one action chunk with shape ``(H, action_dim)``.
|
||||
|
||||
Returns {"obs": dict, "done": bool, "steps": int}.
|
||||
"""
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Release resources. Subclasses may override."""
|
||||
|
||||
# ---- episode-level default implementation ----
|
||||
|
||||
def run_episode(
|
||||
self,
|
||||
ep_id: tuple,
|
||||
predict: Callable[[dict, str, int], np.ndarray],
|
||||
) -> dict:
|
||||
"""Episode loop built from reset, predict, and execute_chunk.
|
||||
|
||||
``predict(observation, instruction, step)`` returns one action chunk.
|
||||
The env remains unaware of the model transport.
|
||||
"""
|
||||
import time
|
||||
|
||||
from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs
|
||||
|
||||
t_ep_start = time.time()
|
||||
initial = self.reset_episode(ep_id)
|
||||
obs = initial["obs"]
|
||||
instruction = initial["instruction"]
|
||||
task_desc = initial.get("task_desc", "")
|
||||
|
||||
max_rounds = self._max_infer_rounds()
|
||||
success = False
|
||||
steps_total = 0
|
||||
for round_idx in range(max_rounds):
|
||||
encoded = encode_raw_obs(obs)
|
||||
chunk = predict(encoded, instruction, round_idx)
|
||||
result = self.execute_chunk(chunk)
|
||||
obs = result["obs"]
|
||||
steps_total += result["steps"]
|
||||
if result["done"]:
|
||||
success = True
|
||||
break
|
||||
|
||||
self.finalize_episode(success)
|
||||
return {
|
||||
"success": bool(success),
|
||||
"steps": steps_total,
|
||||
"elapsed_sec": round(time.time() - t_ep_start, 3),
|
||||
"task_desc": task_desc,
|
||||
}
|
||||
|
||||
def finalize_episode(self, success: bool) -> None:
|
||||
"""Hook for env-specific cleanup after an episode (e.g. save rollouts)."""
|
||||
|
||||
def _max_infer_rounds(self) -> int:
|
||||
"""Return the maximum number of model chunks for one episode."""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} must override _max_infer_rounds when using "
|
||||
"the default run_episode implementation"
|
||||
)
|
||||
@@ -0,0 +1,342 @@
|
||||
"""LIBERO environment implementation.
|
||||
|
||||
The chunk-level API mirrors standard LIBERO rollout semantics:
|
||||
- reset_episode enables rendering, sets initial state, and performs warmup steps.
|
||||
- execute_chunk may skip intermediate image rendering but re-enables rendering
|
||||
before returning an observation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from wall_x._vendor.harrix.envs.base import BaseEnv
|
||||
from wall_x._vendor.harrix.envs.libero_common import (
|
||||
get_rollout_frame,
|
||||
model_action_to_libero_env,
|
||||
save_rollout_video,
|
||||
)
|
||||
from wall_x._vendor.harrix.envs.libero_sim import (
|
||||
create_libero_engine,
|
||||
find_image_observables,
|
||||
get_instruction,
|
||||
get_libero_dummy_action,
|
||||
get_task_suite,
|
||||
load_initial_states,
|
||||
pick_initial_state,
|
||||
resolve_task_info,
|
||||
set_render_enabled,
|
||||
)
|
||||
from wall_x._vendor.harrix.envs.registry import register_env
|
||||
from wall_x._vendor.harrix.eval_config import EvalConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_env("libero")
|
||||
class LiberoEnv(BaseEnv):
|
||||
|
||||
def __init__(self, cfg: EvalConfig, worker_id: int) -> None:
|
||||
self._cfg = cfg
|
||||
self._libero_cfg = cfg.env.libero
|
||||
self._worker_id = worker_id
|
||||
self._seed = cfg.env.seed
|
||||
self._task_suite_name = self._libero_cfg.task_suite_name
|
||||
|
||||
# Task-suite metadata.
|
||||
self._task_suite = get_task_suite(self._task_suite_name)
|
||||
self._num_tasks = self._task_suite.n_tasks
|
||||
self._custom_initial_states = load_initial_states(
|
||||
self._libero_cfg.initial_states_path
|
||||
)
|
||||
|
||||
# Robosuite engine, lazily rebuilt on task changes.
|
||||
self._libero_env = None
|
||||
self._current_task_id: Optional[int] = None
|
||||
self._rebuild_env_per_episode = self._libero_cfg.rebuild_env_per_episode
|
||||
|
||||
# Render-skip state.
|
||||
self._skip_intermediate_render = self._libero_cfg.skip_intermediate_render
|
||||
self._force_render_task_ids = set(
|
||||
self._libero_cfg.force_render_task_indices or []
|
||||
)
|
||||
self._effective_skip_render = self._skip_intermediate_render
|
||||
self._image_obs: list = []
|
||||
self._chunk_granular_render_toggle = (
|
||||
self._libero_cfg.chunk_granular_render_toggle
|
||||
)
|
||||
self._render_enabled_state = False
|
||||
|
||||
# Optional bit-alignment dump for debugging.
|
||||
self._bit_dump_dir = os.environ.get("WALLX_BIT_DUMP_DIR", "").strip() or None
|
||||
if self._bit_dump_dir:
|
||||
os.makedirs(self._bit_dump_dir, exist_ok=True)
|
||||
|
||||
rollout_dir = (self._libero_cfg.rollout_dir or "").strip()
|
||||
if not rollout_dir:
|
||||
rollout_dir = os.environ.get("WALLX_ROLLOUT_DIR", "").strip()
|
||||
if rollout_dir and os.environ.get("WALLX_DISABLE_ROLLOUT", "0") == "1":
|
||||
rollout_dir = ""
|
||||
self._rollout_dir = rollout_dir or None
|
||||
self._rollout_fps = int(self._libero_cfg.rollout_fps)
|
||||
if self._rollout_dir:
|
||||
worker_subdir = f"worker{worker_id}" if cfg.runtime.num_workers > 1 else ""
|
||||
self._rollout_dir = os.path.join(
|
||||
self._rollout_dir,
|
||||
self._task_suite_name,
|
||||
worker_subdir,
|
||||
)
|
||||
os.makedirs(self._rollout_dir, exist_ok=True)
|
||||
logger.info("Rollout MP4 saving enabled: %s", self._rollout_dir)
|
||||
|
||||
# Per-episode state consumed by execute_chunk.
|
||||
self._current_ep: Optional[tuple] = None
|
||||
self._current_task_desc: str = ""
|
||||
self._current_instruction: str = ""
|
||||
self._chunk_counter_in_ep: int = 0
|
||||
self._last_obs_for_dump: Optional[dict] = None
|
||||
self._replay_images: list[np.ndarray] = []
|
||||
self._rollout_saved = False
|
||||
|
||||
# ---- BaseEnv API ----
|
||||
|
||||
@classmethod
|
||||
def enumerate_episodes(cls, cfg: EvalConfig) -> list[tuple]:
|
||||
libero_cfg = cfg.env.libero
|
||||
suite = libero_cfg.task_suite_name
|
||||
|
||||
if libero_cfg.task_indices is not None:
|
||||
task_indices = [int(x) for x in libero_cfg.task_indices]
|
||||
else:
|
||||
ts = get_task_suite(suite)
|
||||
task_indices = list(range(ts.n_tasks))
|
||||
|
||||
eps = []
|
||||
for tid in task_indices:
|
||||
for epi in range(libero_cfg.num_trials_per_task):
|
||||
eps.append((suite, tid, epi))
|
||||
return eps
|
||||
|
||||
@property
|
||||
def robot_spec(self) -> dict:
|
||||
"""Return robot spec for driver/adapter validation."""
|
||||
from wall_x._vendor.harrix.utils.train_config import (
|
||||
load_train_config_with_ckpt_overlay,
|
||||
)
|
||||
|
||||
train_cfg = load_train_config_with_ckpt_overlay(
|
||||
self._cfg.model.train_config_path,
|
||||
self._cfg.model.checkpoint_path,
|
||||
)
|
||||
return {
|
||||
"dof_layout": train_cfg.get("dof_config", {}),
|
||||
"cam_names": list(self._cfg.model.cam_names),
|
||||
"norm_key": self._cfg.model.norm_key,
|
||||
}
|
||||
|
||||
def _max_infer_rounds(self) -> int:
|
||||
return self._libero_cfg.max_infer_times
|
||||
|
||||
def reset_episode(self, ep_id: tuple) -> dict:
|
||||
suite, task_id, ep_idx = ep_id
|
||||
if suite != self._task_suite_name:
|
||||
raise ValueError(
|
||||
f"env bound to suite={self._task_suite_name}, got ep with suite={suite}"
|
||||
)
|
||||
|
||||
need_rebuild = self._rebuild_env_per_episode or self._current_task_id != task_id
|
||||
if need_rebuild:
|
||||
self._rebuild_env(task_id)
|
||||
|
||||
task_desc, default_states = resolve_task_info(self._task_suite, task_id)
|
||||
init_state = pick_initial_state(
|
||||
self._libero_cfg.initial_states_path,
|
||||
self._custom_initial_states,
|
||||
task_desc,
|
||||
default_states,
|
||||
ep_idx,
|
||||
)
|
||||
|
||||
if not need_rebuild:
|
||||
self._libero_env.reset()
|
||||
obs = self._libero_env.set_init_state(init_state)
|
||||
if obs is None:
|
||||
raise RuntimeError("set_init_state returned None")
|
||||
|
||||
set_render_enabled(self._image_obs, True)
|
||||
self._render_enabled_state = True
|
||||
|
||||
dummy_action = get_libero_dummy_action()
|
||||
for _ in range(10):
|
||||
obs, _, _, _ = self._libero_env.step(dummy_action)
|
||||
|
||||
self._current_ep = (task_id, ep_idx)
|
||||
self._current_task_desc = task_desc
|
||||
self._current_instruction = get_instruction(task_desc)
|
||||
self._chunk_counter_in_ep = 0
|
||||
self._last_obs_for_dump = obs
|
||||
self._begin_rollout_capture(obs)
|
||||
|
||||
return {
|
||||
"obs": obs,
|
||||
"instruction": self._current_instruction,
|
||||
"task_desc": task_desc,
|
||||
}
|
||||
|
||||
def execute_chunk(self, actions: np.ndarray) -> dict:
|
||||
actions = np.asarray(actions, dtype=np.float32)
|
||||
H = actions.shape[0]
|
||||
|
||||
if self._bit_dump_dir and self._current_ep is not None:
|
||||
self._dump_chunk_npz(
|
||||
(self._task_suite_name, *self._current_ep),
|
||||
self._chunk_counter_in_ep,
|
||||
self._last_obs_for_dump or {},
|
||||
actions,
|
||||
)
|
||||
self._chunk_counter_in_ep += 1
|
||||
|
||||
skip_render = self._chunk_skip_render()
|
||||
# Disable rendering at the chunk start when render-skip is enabled.
|
||||
if skip_render and self._image_obs:
|
||||
if self._render_enabled_state:
|
||||
set_render_enabled(self._image_obs, False)
|
||||
self._render_enabled_state = False
|
||||
|
||||
last_obs = None
|
||||
done = False
|
||||
steps = 0
|
||||
for step_idx in range(H):
|
||||
# Re-enable rendering before the final step to return a fresh image.
|
||||
if (
|
||||
skip_render
|
||||
and self._image_obs
|
||||
and step_idx == H - 1
|
||||
and not self._render_enabled_state
|
||||
):
|
||||
set_render_enabled(self._image_obs, True)
|
||||
self._render_enabled_state = True
|
||||
|
||||
action = model_action_to_libero_env(actions[step_idx].reshape(-1))
|
||||
obs, _, done_flag, _ = self._libero_env.step(action)
|
||||
last_obs = obs
|
||||
steps += 1
|
||||
self._append_rollout_frame(obs)
|
||||
if bool(done_flag):
|
||||
done = True
|
||||
# Do not add an extra simulator step on early success; just make
|
||||
# sure future rendering is enabled.
|
||||
if skip_render and self._image_obs and not self._render_enabled_state:
|
||||
set_render_enabled(self._image_obs, True)
|
||||
self._render_enabled_state = True
|
||||
break
|
||||
|
||||
self._last_obs_for_dump = last_obs
|
||||
return {"obs": last_obs, "done": done, "steps": steps}
|
||||
|
||||
def finalize_episode(self, success: bool) -> None:
|
||||
self._save_episode_rollout(success)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
if self._libero_env is not None:
|
||||
try:
|
||||
self._libero_env.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._libero_env = None
|
||||
|
||||
# ---- internals ----
|
||||
|
||||
def _chunk_skip_render(self) -> bool:
|
||||
"""Skip intermediate renders unless rollout MP4 saving needs every frame."""
|
||||
return self._effective_skip_render and self._rollout_dir is None
|
||||
|
||||
def _begin_rollout_capture(self, obs: dict | None) -> None:
|
||||
self._replay_images = []
|
||||
self._rollout_saved = False
|
||||
if self._rollout_dir and obs is not None:
|
||||
if self._image_obs and not self._render_enabled_state:
|
||||
set_render_enabled(self._image_obs, True)
|
||||
self._render_enabled_state = True
|
||||
self._replay_images.append(get_rollout_frame(obs))
|
||||
|
||||
def _append_rollout_frame(self, obs: dict | None) -> None:
|
||||
if self._rollout_dir and obs is not None:
|
||||
self._replay_images.append(get_rollout_frame(obs))
|
||||
|
||||
def _save_episode_rollout(self, success: bool) -> None:
|
||||
if (
|
||||
not self._rollout_dir
|
||||
or self._rollout_saved
|
||||
or not self._replay_images
|
||||
or self._current_ep is None
|
||||
):
|
||||
return
|
||||
task_id, ep_idx = self._current_ep
|
||||
try:
|
||||
mp4_path = save_rollout_video(
|
||||
self._rollout_dir,
|
||||
self._replay_images,
|
||||
task_id=task_id,
|
||||
episode_idx=ep_idx,
|
||||
success=success,
|
||||
task_description=self._current_task_desc,
|
||||
fps=self._rollout_fps,
|
||||
)
|
||||
self._rollout_saved = True
|
||||
if mp4_path:
|
||||
logger.info("Saved rollout MP4: %s", mp4_path)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to save rollout MP4 for task%d ep%d: %s",
|
||||
task_id,
|
||||
ep_idx,
|
||||
exc,
|
||||
)
|
||||
|
||||
def _rebuild_env(self, task_id: int) -> None:
|
||||
if self._libero_env is not None:
|
||||
try:
|
||||
self._libero_env.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._libero_env = None
|
||||
gc.collect()
|
||||
|
||||
self._libero_env = create_libero_engine(
|
||||
task_id=task_id,
|
||||
task_suite_name=self._task_suite_name,
|
||||
resolution=256,
|
||||
seed=self._seed,
|
||||
)
|
||||
self._current_task_id = task_id
|
||||
self._image_obs = find_image_observables(self._libero_env)
|
||||
self._effective_skip_render = (
|
||||
self._skip_intermediate_render
|
||||
and task_id not in self._force_render_task_ids
|
||||
)
|
||||
|
||||
def _dump_chunk_npz(
|
||||
self, ep_id, chunk_idx: int, raw_obs: dict, chunk_actions: np.ndarray
|
||||
) -> None:
|
||||
_, task_id, ep_idx = ep_id
|
||||
path = os.path.join(
|
||||
self._bit_dump_dir, f"t{task_id}_ep{ep_idx}_c{chunk_idx}.npz"
|
||||
)
|
||||
fields = {"action_chunk": np.asarray(chunk_actions, dtype=np.float32)}
|
||||
for k in (
|
||||
"robot0_eef_pos",
|
||||
"robot0_eef_quat",
|
||||
"robot0_gripper_qpos",
|
||||
"agentview_image",
|
||||
"robot0_eye_in_hand_image",
|
||||
):
|
||||
v = raw_obs.get(k)
|
||||
if v is not None:
|
||||
fields[k] = np.asarray(v)
|
||||
np.savez_compressed(path, **fields)
|
||||
@@ -0,0 +1,359 @@
|
||||
"""Env-adapter IO helpers for single-arm LIBERO tasks.
|
||||
|
||||
Env code extracts a compact ndarray payload from raw LIBERO observations.
|
||||
Adapters then build proprioception, masks, and 7-dof right-arm action chunks
|
||||
from that payload. This module intentionally depends only on NumPy so adapter
|
||||
processes can import it without importing robosuite or LIBERO.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
# Fallback single-arm LIBERO dof_config when train_config does not provide one.
|
||||
_LIBERO_FALLBACK_DOF_CONFIG = {
|
||||
"follow_right_ee_cartesian_pos": 3,
|
||||
"follow_right_ee_rotation": 3,
|
||||
"follow_right_gripper": 1,
|
||||
}
|
||||
|
||||
_LIBERO_FALLBACK_AGENT_POS_CONFIG = dict(_LIBERO_FALLBACK_DOF_CONFIG)
|
||||
_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",))
|
||||
|
||||
|
||||
def _resolve_dof_config(train_config: dict) -> dict:
|
||||
return (
|
||||
train_config.get("dof_config")
|
||||
or train_config.get("task", {}).get("dof_config")
|
||||
or _LIBERO_FALLBACK_DOF_CONFIG
|
||||
)
|
||||
|
||||
|
||||
def _resolve_agent_pos_config(train_config: dict) -> dict:
|
||||
return (
|
||||
train_config.get("agent_pos_config")
|
||||
or train_config.get("task", {}).get("agent_pos_config")
|
||||
or _LIBERO_FALLBACK_AGENT_POS_CONFIG
|
||||
)
|
||||
|
||||
|
||||
def _move_virtual_keys_to_tail(layout: dict) -> dict:
|
||||
head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS}
|
||||
tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS}
|
||||
return {**head, **tail}
|
||||
|
||||
|
||||
def _effective_agent_pos_config(train_config: dict, state_values: dict) -> dict:
|
||||
config = _move_virtual_keys_to_tail(dict(_resolve_agent_pos_config(train_config)))
|
||||
gripper_key = next(
|
||||
(
|
||||
key
|
||||
for key in config
|
||||
if key.replace("follow_", "").replace("master_", "") == "right_gripper"
|
||||
),
|
||||
None,
|
||||
)
|
||||
if gripper_key is None:
|
||||
return config
|
||||
|
||||
old_dim = int(config[gripper_key])
|
||||
target_dim = None
|
||||
override = os.environ.get("WALLX_LIBERO_STATE_GRIPPER_DIM")
|
||||
if override:
|
||||
target_dim = int(override)
|
||||
elif os.environ.get("WALLX_LIBERO_AUTO_STATE_GRIPPER_DIM", "1") != "0":
|
||||
norm_dim = int(train_config.get("_libero_proprio_norm_dim") or 0)
|
||||
real_dim = sum(v for k, v in config.items() if k not in _VIRTUAL_TAIL_KEYS)
|
||||
if norm_dim == real_dim + 1:
|
||||
target_dim = old_dim + 1
|
||||
|
||||
available_dim = state_values["right_gripper"].shape[1]
|
||||
if target_dim is None or target_dim == old_dim or target_dim > available_dim:
|
||||
return config
|
||||
|
||||
config[gripper_key] = target_dim
|
||||
delta = target_dim - old_dim
|
||||
if "action_padding" in config:
|
||||
config["action_padding"] = max(0, int(config["action_padding"]) - delta)
|
||||
return config
|
||||
|
||||
|
||||
def _build_right_arm_state_values(obs_ndarrays: dict) -> dict[str, np.ndarray]:
|
||||
"""Bare-key state tensors for single-arm LIBERO proprio construction."""
|
||||
rot3 = np.asarray(obs_ndarrays["eef_axisangle"], dtype=np.float32).reshape(1, 3)
|
||||
values: dict[str, np.ndarray] = {
|
||||
"right_ee_cartesian_pos": np.asarray(
|
||||
obs_ndarrays["eef_pos"], dtype=np.float32
|
||||
).reshape(1, 3),
|
||||
"right_ee_rotation": rot3,
|
||||
"right_gripper": np.asarray(obs_ndarrays["gripper"], dtype=np.float32).reshape(
|
||||
1, -1
|
||||
),
|
||||
}
|
||||
from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb
|
||||
|
||||
rot6d = euler_to_matrix_zyx_6d_nb(rot3.astype(np.float64)).reshape(1, 6)
|
||||
values["right_ee_rotation_6D"] = rot6d.astype(np.float32)
|
||||
return values
|
||||
|
||||
|
||||
# Auxiliary action keys that should be masked out for single-arm LIBERO.
|
||||
_DOF_MASK_ZERO_KEYS = frozenset(
|
||||
(
|
||||
"follow_left_ee_cartesian_pos",
|
||||
"follow_left_ee_rotation",
|
||||
"follow_left_ee_rotation_6D",
|
||||
"follow_left_gripper",
|
||||
"head_actions",
|
||||
"height",
|
||||
"velocity_decomposed",
|
||||
"action_padding",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# LIBERO raw observation decoding helpers.
|
||||
# ============================================================
|
||||
|
||||
|
||||
def _get_libero_image(obs: dict) -> np.ndarray:
|
||||
"""Return the third-person camera image, rotated to match preprocessing."""
|
||||
return obs["agentview_image"][::-1, ::-1]
|
||||
|
||||
|
||||
def get_rollout_frame(obs: dict) -> np.ndarray:
|
||||
"""Return one RGB frame for rollout MP4 saving."""
|
||||
return np.asarray(_get_libero_image(obs), dtype=np.uint8)
|
||||
|
||||
|
||||
def _get_libero_wrist_image(obs: dict) -> np.ndarray:
|
||||
"""Return the wrist camera image, rotated to match preprocessing."""
|
||||
return obs["robot0_eye_in_hand_image"][::-1, ::-1]
|
||||
|
||||
|
||||
def _quat2axisangle(quat) -> np.ndarray:
|
||||
"""Convert an xyzw quaternion to a 3D axis-angle vector."""
|
||||
if quat[3] > 1.0:
|
||||
quat[3] = 1.0
|
||||
elif quat[3] < -1.0:
|
||||
quat[3] = -1.0
|
||||
|
||||
den = np.sqrt(1.0 - quat[3] * quat[3])
|
||||
if math.isclose(den, 0.0):
|
||||
return np.zeros(3)
|
||||
|
||||
return (quat[:3] * 2.0 * math.acos(quat[3])) / den
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Env-side to adapter-side observation encoding.
|
||||
# ============================================================
|
||||
|
||||
|
||||
def encode_raw_obs(raw_obs: dict) -> dict:
|
||||
"""Extract the minimal ndarray payload from a raw LIBERO observation.
|
||||
|
||||
The payload contains three 1-D state arrays and two rotated image arrays.
|
||||
"""
|
||||
if "agentview_image" not in raw_obs:
|
||||
raise KeyError(
|
||||
"agentview_image missing in raw_obs; render-skip may have returned "
|
||||
"a stale observation"
|
||||
)
|
||||
return {
|
||||
"eef_pos": np.asarray(raw_obs["robot0_eef_pos"], dtype=np.float32),
|
||||
"eef_axisangle": np.asarray(
|
||||
_quat2axisangle(raw_obs["robot0_eef_quat"]), dtype=np.float32
|
||||
),
|
||||
"gripper": np.asarray(raw_obs["robot0_gripper_qpos"], dtype=np.float32),
|
||||
"face_view": _get_libero_image(raw_obs),
|
||||
"wrist_view": _get_libero_wrist_image(raw_obs),
|
||||
}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Adapter-side proprioception and mask construction.
|
||||
# ============================================================
|
||||
|
||||
|
||||
def encode_proprio(
|
||||
obs_ndarrays: dict,
|
||||
train_config: dict,
|
||||
action_horizon: int,
|
||||
) -> dict:
|
||||
"""Convert a single-arm LIBERO ndarray payload into model input fields.
|
||||
|
||||
Returned fields include proprioception, agent_pos_mask, dof_mask,
|
||||
face_view, and right_wrist_view.
|
||||
"""
|
||||
state_values = _build_right_arm_state_values(obs_ndarrays)
|
||||
agent_pos_config = _effective_agent_pos_config(train_config, state_values)
|
||||
dof_config = _move_virtual_keys_to_tail(dict(_resolve_dof_config(train_config)))
|
||||
|
||||
propri_parts: list[np.ndarray] = []
|
||||
mask_parts: list[np.ndarray] = []
|
||||
for key, dim in agent_pos_config.items():
|
||||
bare = key.replace("follow_", "").replace("master_", "")
|
||||
if bare in state_values:
|
||||
v = state_values[bare]
|
||||
if bare == "right_gripper" and v.shape[1] > dim:
|
||||
v = v[:, :dim]
|
||||
if v.shape[1] != dim:
|
||||
raise ValueError(
|
||||
f"agent_pos_config[{key!r}]={dim} does not match "
|
||||
f"observation dimension {v.shape[1]}"
|
||||
)
|
||||
propri_parts.append(v)
|
||||
mask_parts.append(np.ones((1, dim), dtype=np.float32))
|
||||
else:
|
||||
propri_parts.append(np.zeros((1, dim), dtype=np.float32))
|
||||
mask_parts.append(np.zeros((1, dim), dtype=np.float32))
|
||||
|
||||
# (1, 1, D)
|
||||
proprioception = np.concatenate(propri_parts, axis=1)[None]
|
||||
agent_pos_mask = np.concatenate(mask_parts, axis=1)[None]
|
||||
|
||||
# dof_mask: (1, T, D_action)
|
||||
total_dof = sum(dof_config.values())
|
||||
dof_mask = np.ones((1, action_horizon, total_dof))
|
||||
start = 0
|
||||
for key, dim in dof_config.items():
|
||||
if key in _DOF_MASK_ZERO_KEYS:
|
||||
dof_mask[:, :, start : start + dim] = 0
|
||||
start += dim
|
||||
|
||||
return {
|
||||
"proprioception": proprioception.astype(np.float32),
|
||||
"agent_pos_mask": agent_pos_mask.astype(np.float32),
|
||||
"dof_mask": dof_mask,
|
||||
"face_view": obs_ndarrays["face_view"],
|
||||
"right_wrist_view": obs_ndarrays["wrist_view"],
|
||||
}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Adapter-side action decoding.
|
||||
# ============================================================
|
||||
|
||||
|
||||
def decode_chunk(predict_action: np.ndarray, train_config: dict) -> np.ndarray:
|
||||
"""Extract a 7-dof right-arm chunk from model action output."""
|
||||
if predict_action.ndim == 3:
|
||||
predict_action = predict_action[0]
|
||||
|
||||
dof_config = _resolve_dof_config(train_config)
|
||||
slices: dict[str, slice] = {}
|
||||
start = 0
|
||||
for key, dim in dof_config.items():
|
||||
bare = key.replace("follow_", "").replace("master_", "")
|
||||
slices[bare] = slice(start, start + dim)
|
||||
start += dim
|
||||
|
||||
pos = predict_action[:, slices["right_ee_cartesian_pos"]]
|
||||
grip = predict_action[:, slices["right_gripper"]]
|
||||
|
||||
if "right_ee_rotation_6D" in slices:
|
||||
from wall_x._vendor.x2robot_utils.geometry import so3_to_euler_zyx_batch_nb
|
||||
|
||||
rot6d = np.asarray(
|
||||
predict_action[:, slices["right_ee_rotation_6D"]], dtype=np.float64
|
||||
)
|
||||
rot = so3_to_euler_zyx_batch_nb(rot6d).astype(np.float32)
|
||||
elif "right_ee_rotation" in slices:
|
||||
rot = predict_action[:, slices["right_ee_rotation"]]
|
||||
else:
|
||||
raise KeyError(
|
||||
"dof_config has no right-arm rotation slice "
|
||||
f"(keys={list(slices.keys())})"
|
||||
)
|
||||
|
||||
return np.concatenate([pos, rot, grip], axis=1)
|
||||
|
||||
|
||||
def gripper_model_to_libero_osc(grip_2d: np.ndarray) -> float:
|
||||
"""Map model gripper output to robosuite OSC_POSE gripper command in [-1, 1]."""
|
||||
g = np.asarray(grip_2d, dtype=np.float64).reshape(-1)
|
||||
if g.size == 0:
|
||||
raise ValueError("empty gripper action")
|
||||
cmd = float(g[0])
|
||||
if os.environ.get("WALLX_LIBERO_GRIPPER_BINARIZE", "0") == "1":
|
||||
if abs(cmd) < 1e-6:
|
||||
cmd = -1.0
|
||||
else:
|
||||
cmd = float(np.sign(cmd))
|
||||
if os.environ.get("WALLX_LIBERO_INVERT_GRIPPER", "0") == "1":
|
||||
cmd *= -1.0
|
||||
return cmd
|
||||
|
||||
|
||||
def _sanitize_task_description(task_description: str, max_len: int = 50) -> str:
|
||||
return (
|
||||
task_description.lower()
|
||||
.replace(" ", "_")
|
||||
.replace("\n", "_")
|
||||
.replace(".", "_")[:max_len]
|
||||
)
|
||||
|
||||
|
||||
def save_rollout_video(
|
||||
rollout_dir: str,
|
||||
rollout_images: list[np.ndarray],
|
||||
*,
|
||||
task_id: int,
|
||||
episode_idx: int,
|
||||
success: bool,
|
||||
task_description: str,
|
||||
fps: int = 30,
|
||||
) -> str | None:
|
||||
"""Save an MP4 replay of one LIBERO episode."""
|
||||
if not rollout_images:
|
||||
return None
|
||||
|
||||
import imageio
|
||||
|
||||
os.makedirs(rollout_dir, exist_ok=True)
|
||||
task_slug = _sanitize_task_description(task_description)
|
||||
mp4_path = os.path.join(
|
||||
rollout_dir,
|
||||
f"task{task_id}_ep{episode_idx}--success={int(success)}--{task_slug}.mp4",
|
||||
)
|
||||
writer = imageio.get_writer(mp4_path, fps=fps, macro_block_size=1)
|
||||
try:
|
||||
for img in rollout_images:
|
||||
writer.append_data(np.asarray(img, dtype=np.uint8))
|
||||
finally:
|
||||
writer.close()
|
||||
return mp4_path
|
||||
|
||||
|
||||
def model_action_to_libero_env(action: np.ndarray) -> np.ndarray:
|
||||
"""Convert model output to the robosuite OSC_POSE 7D action.
|
||||
|
||||
Internal LIBERO evaluation passes the model's 7D chunk directly to
|
||||
``env.step``. Keep that as the public default; conversion modes remain
|
||||
available only for ablations through environment variables.
|
||||
"""
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
a = np.asarray(action, dtype=np.float64).reshape(-1)
|
||||
if a.size not in (7, 8):
|
||||
raise ValueError(f"expected 7D or 8D model action, got shape {a.shape}")
|
||||
|
||||
pos_delta = a[:3]
|
||||
rot_mode = os.environ.get("WALLX_LIBERO_ROT_MODE", "direct").strip()
|
||||
if rot_mode == "euler_zyx_to_rotvec":
|
||||
rot_aa = R.from_euler("zyx", a[3:6]).as_rotvec()
|
||||
elif rot_mode == "direct":
|
||||
rot_aa = a[3:6]
|
||||
else:
|
||||
raise ValueError(
|
||||
"WALLX_LIBERO_ROT_MODE must be 'euler_zyx_to_rotvec' or 'direct', "
|
||||
f"got {rot_mode!r}"
|
||||
)
|
||||
grip = gripper_model_to_libero_osc(a[6:8] if a.size == 8 else a[6:7])
|
||||
return np.concatenate([pos_delta, rot_aa, np.array([grip], dtype=np.float64)])
|
||||
@@ -0,0 +1,181 @@
|
||||
"""LIBERO benchmark and robosuite engine helpers.
|
||||
|
||||
This module may import robosuite and LIBERO. Adapter-side code should use
|
||||
``libero_common.py`` instead, which only depends on NumPy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
import numpy as np
|
||||
from robosuite.wrappers import VisualizationWrapper
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# One-time side effect: auto-create ~/.libero/config.yaml so importing LIBERO
|
||||
# does not trigger an interactive prompt.
|
||||
# ============================================================
|
||||
|
||||
|
||||
def _ensure_libero_config() -> None:
|
||||
import yaml as _yaml
|
||||
|
||||
libero_config_path = os.environ.get(
|
||||
"LIBERO_CONFIG_PATH", os.path.expanduser("~/.libero")
|
||||
)
|
||||
config_file = os.path.join(libero_config_path, "config.yaml")
|
||||
|
||||
if not os.path.exists(config_file):
|
||||
os.makedirs(libero_config_path, exist_ok=True)
|
||||
import libero.libero as _libero_pkg
|
||||
|
||||
benchmark_root = os.path.dirname(os.path.abspath(_libero_pkg.__file__))
|
||||
default_paths = {
|
||||
"benchmark_root": benchmark_root,
|
||||
"bddl_files": os.path.join(benchmark_root, "./bddl_files"),
|
||||
"init_states": os.path.join(benchmark_root, "./init_files"),
|
||||
"datasets": os.path.join(benchmark_root, "../datasets"),
|
||||
"assets": os.path.join(benchmark_root, "./assets"),
|
||||
}
|
||||
with open(config_file, "w") as f:
|
||||
_yaml.dump(default_paths, f)
|
||||
logger.info("Auto-created LIBERO config: %s", config_file)
|
||||
|
||||
|
||||
_ensure_libero_config()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# task-suite entry point
|
||||
# ============================================================
|
||||
|
||||
|
||||
def get_task_suite(task_suite_name: str):
|
||||
"""Load a LIBERO task suite."""
|
||||
from libero.libero import benchmark
|
||||
|
||||
return benchmark.get_benchmark_dict()[task_suite_name]()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# actions
|
||||
# ============================================================
|
||||
|
||||
|
||||
def get_libero_dummy_action() -> list[float]:
|
||||
"""Return the 7-dof dummy action used for episode warmup."""
|
||||
return [0, 0, 0, 0, 0, 0, -1]
|
||||
|
||||
|
||||
# ============================================================
|
||||
# robosuite engine factory
|
||||
# ============================================================
|
||||
|
||||
|
||||
def create_libero_engine(
|
||||
task_id: int,
|
||||
task_suite_name: str,
|
||||
resolution: int = 256,
|
||||
seed: int = 7,
|
||||
) -> Any:
|
||||
"""Construct one LIBERO robosuite engine."""
|
||||
from libero.libero import get_libero_path
|
||||
from libero.libero.envs import OffScreenRenderEnv
|
||||
|
||||
task_suite = get_task_suite(task_suite_name)
|
||||
task = task_suite.get_task(task_id)
|
||||
task_bddl_file = os.path.join(
|
||||
get_libero_path("bddl_files"), task.problem_folder, task.bddl_file
|
||||
)
|
||||
env = OffScreenRenderEnv(
|
||||
bddl_file_name=task_bddl_file,
|
||||
camera_heights=resolution,
|
||||
camera_widths=resolution,
|
||||
)
|
||||
# The seed still affects object poses even when an initial state is fixed.
|
||||
env.seed(seed)
|
||||
env.env = VisualizationWrapper(env.env)
|
||||
env.env.set_visualization_setting(setting="grippers", visible=False)
|
||||
return env
|
||||
|
||||
|
||||
# ============================================================
|
||||
# task metadata / initial states
|
||||
# ============================================================
|
||||
|
||||
|
||||
def load_initial_states(initial_states_path: str) -> Optional[dict]:
|
||||
"""Load custom initial states, or return None for suite defaults."""
|
||||
if initial_states_path == "DEFAULT":
|
||||
return None
|
||||
with open(initial_states_path, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def resolve_task_info(task_suite, task_id: int) -> tuple[str, Any]:
|
||||
"""Return ``(task_desc, default_initial_states)`` for one task id."""
|
||||
num_tasks = task_suite.n_tasks
|
||||
if task_id < 0 or task_id >= num_tasks:
|
||||
raise ValueError(f"invalid task_id={task_id}, num_tasks={num_tasks}")
|
||||
task = task_suite.get_task(task_id)
|
||||
return task.language, task_suite.get_task_init_states(task_id)
|
||||
|
||||
|
||||
def pick_initial_state(
|
||||
initial_states_path: str,
|
||||
custom_initial_states: Optional[dict],
|
||||
task_desc: str,
|
||||
default_states: Any,
|
||||
episode_idx: int,
|
||||
) -> np.ndarray:
|
||||
"""Pick one initial state from suite defaults or a custom states file."""
|
||||
if initial_states_path == "DEFAULT":
|
||||
if default_states is None:
|
||||
raise ValueError("default states missing for DEFAULT mode")
|
||||
return default_states[episode_idx]
|
||||
|
||||
if custom_initial_states is None:
|
||||
raise ValueError(f"custom initial states not loaded for {initial_states_path}")
|
||||
key = task_desc.replace(" ", "_")
|
||||
ep_key = f"demo_{episode_idx}"
|
||||
record = custom_initial_states[key][ep_key]
|
||||
if not record["success"]:
|
||||
raise ValueError(f"expert demo failed for {ep_key}")
|
||||
return np.array(record["initial_state"])
|
||||
|
||||
|
||||
def get_instruction(task_desc: str) -> str:
|
||||
"""Return the instruction text for a LIBERO task description."""
|
||||
return task_desc
|
||||
|
||||
|
||||
# ============================================================
|
||||
# render-skip: directly assign obs._enabled to avoid set_enabled() side effects.
|
||||
# ============================================================
|
||||
|
||||
|
||||
def find_image_observables(env) -> list:
|
||||
"""Find image observables along the env.env wrapper chain."""
|
||||
cur = env
|
||||
seen = set()
|
||||
while cur is not None and id(cur) not in seen:
|
||||
seen.add(id(cur))
|
||||
if hasattr(cur, "_observables") and isinstance(cur._observables, dict):
|
||||
return [
|
||||
obs
|
||||
for obs in cur._observables.values()
|
||||
if getattr(obs, "modality", None) == "image"
|
||||
]
|
||||
cur = getattr(cur, "env", None)
|
||||
return []
|
||||
|
||||
|
||||
def set_render_enabled(image_obs_list, enabled: bool) -> None:
|
||||
for obs in image_obs_list:
|
||||
obs._enabled = enabled
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Environment registry and factory."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Type
|
||||
|
||||
from wall_x._vendor.harrix.eval_config import EvalConfig
|
||||
from wall_x._vendor.harrix.envs.base import BaseEnv
|
||||
|
||||
|
||||
_REGISTRY: dict[str, Type[BaseEnv]] = {}
|
||||
|
||||
|
||||
def register_env(name: str):
|
||||
def deco(cls: Type[BaseEnv]):
|
||||
if name in _REGISTRY:
|
||||
raise ValueError(
|
||||
f"env {name!r} already registered (cls={_REGISTRY[name].__name__})"
|
||||
)
|
||||
_REGISTRY[name] = cls
|
||||
return cls
|
||||
|
||||
return deco
|
||||
|
||||
|
||||
def _get_class(cfg: EvalConfig) -> Type[BaseEnv]:
|
||||
t = cfg.env.type
|
||||
cls = _REGISTRY.get(t)
|
||||
if cls is None:
|
||||
if t == "libero":
|
||||
import wall_x._vendor.harrix.envs as _envs
|
||||
|
||||
exc = getattr(_envs, "_LIBERO_IMPORT_ERROR", None)
|
||||
if exc is not None:
|
||||
raise RuntimeError(
|
||||
"LIBERO evaluation dependencies are not installed. "
|
||||
"Install LIBERO/robosuite and their simulator dependencies "
|
||||
"before using env.type='libero'."
|
||||
) from exc
|
||||
raise ValueError(f"unknown env type={t!r}, registered: {sorted(_REGISTRY)}")
|
||||
return cls
|
||||
|
||||
|
||||
def build_env(cfg: EvalConfig, worker_id: int) -> BaseEnv:
|
||||
return _get_class(cfg)(cfg, worker_id)
|
||||
|
||||
|
||||
def enumerate_episodes_for(cfg: EvalConfig) -> list[tuple]:
|
||||
return _get_class(cfg).enumerate_episodes(cfg)
|
||||
|
||||
|
||||
def registered_envs() -> list[str]:
|
||||
return sorted(_REGISTRY)
|
||||
Reference in New Issue
Block a user