admin/any4lerobot
1
0
Fork 0
mirror of https://github.com/Tavish9/any4lerobot.git synced 2026-07-24 12:35:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor AgiBot converter onto generic pipeline

Browse Source 
Co-authored-by: Codex <codex@openai.com>
This commit is contained in:
Tavish
2026-06-22 13:12:58 +08:00
parent 2baee72741
commit 8afce7fa23
3 changed files with 414 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agibot2lerobot/agibot_utils/agibot_utils.py
+70 -18
View File
@@ -6,22 +6,67 @@ import numpy as np
from PIL import Image
def get_task_info(task_json_path: str) -> dict:
def get_task_info(task_json_path: str | Path) -> list[dict]:
with open(task_json_path, "r") as f:
task_info: list = json.load(f)
task_info.sort(key=lambda episode: episode["episode_id"])
return task_info
def get_task_id(task_json_path: str | Path) -> str:
return Path(task_json_path).stem.split("_")[-1]
def get_episode_ids(src_path: str | Path, task_id: str | int) -> list[int]:
observations_dir = Path(src_path) / "observations" / str(task_id)
return sorted(
int(path.name)
for path in observations_dir.glob("*")
if path.is_dir() and path.name.isdigit()
)
def get_episode_videos(
src_path: str | Path,
task_id: str | int,
episode_id: int,
agibot_world_config: dict,
) -> dict[str, Path]:
ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}"
return {
f"observation.images.{key}": ob_dir / "videos" / f"{key}_color.mp4"
if "sensor" not in key
else ob_dir / "tactile" / f"{key}.mp4" # HACK: handle tactile videos
for key in agibot_world_config["images"]
if "depth" not in key
}
def has_episode_videos(
src_path: str | Path,
task_id: str | int,
episode_id: int,
agibot_world_config: dict,
) -> bool:
videos = get_episode_videos(src_path, task_id, episode_id, agibot_world_config)
return all(video_path.exists() for video_path in videos.values())
def load_depths(root_dir: str, camera_name: str):
cam_path = Path(root_dir)
all_imgs = sorted(list(cam_path.glob(f"{camera_name}*")))
return [np.array(Image.open(f)).astype(np.float32)[:, :, None] / 1000 for f in all_imgs]
return [
np.array(Image.open(f)).astype(np.float32)[:, :, None] / 1000 for f in all_imgs
]
def load_local_dataset(
episode_id: int, src_path: str, task_id: int, save_depth: bool, AgiBotWorld_CONFIG: dict
) -> tuple[list, dict]:
episode_id: int,
src_path: str | Path,
task_id: str | int,
save_depth: bool,
AgiBotWorld_CONFIG: dict,
) -> tuple[int, list[dict], dict[str, Path]] | None:
"""Load local dataset and return a dict with observations and actions"""
ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}"
proprio_dir = Path(src_path) / f"proprio_stats/{task_id}/{episode_id}"
@@ -30,9 +75,13 @@ def load_local_dataset(
action = {}
with h5py.File(proprio_dir / "proprio_stats.h5", "r") as f:
for key in AgiBotWorld_CONFIG["states"]:
state[f"observation.states.{key}"] = np.array(f["state/" + key.replace(".", "/")], dtype=np.float32)
state[f"observation.states.{key}"] = np.array(
f["state/" + key.replace(".", "/")], dtype=np.float32
)
for key in AgiBotWorld_CONFIG["actions"]:
action[f"actions.{key}"] = np.array(f["action/" + key.replace(".", "/")], dtype=np.float32)
action[f"actions.{key}"] = np.array(
f["action/" + key.replace(".", "/")], dtype=np.float32
)
# HACK: agibot team forgot to pad or filter some of the values
num_frames = len(next(iter(state.values())))
@@ -42,7 +91,10 @@ def load_local_dataset(
elif len(action_value) < num_frames:
state_key = action_key.replace("actions", "state").replace(".", "/")
new_action_value = np.array(f[state_key], dtype=np.float32).copy()
action_index_key = "/".join(list(action_key.replace("actions", "action").split(".")[:-1]) + ["index"])
action_index_key = "/".join(
list(action_key.replace("actions", "action").split(".")[:-1])
+ ["index"]
)
action_index = np.array(f[action_index_key])
# agibot lost end index, replace it with joint
if not action_index.size:
@@ -52,11 +104,13 @@ def load_local_dataset(
action[action_key] = new_action_value
elif len(action_value) > num_frames:
print("corrupt data, skipping")
return episode_id, [], {"dummy_video": Path("/path/to/no_exist")}
return None
if save_depth:
depth_imgs = load_depths(ob_dir / "depth", "head_depth")
assert num_frames == len(depth_imgs), "Number of images and states are not equal"
assert num_frames == len(depth_imgs), (
"Number of images and states are not equal"
)
state_key_prefix_len = len("observation.states.")
action_key_prefix_len = len("actions.")
@@ -68,7 +122,9 @@ def load_local_dataset(
if value.size
else np.zeros(
AgiBotWorld_CONFIG["states"][key[state_key_prefix_len:]]["shape"],
dtype=AgiBotWorld_CONFIG["states"][key[state_key_prefix_len:]]["dtype"],
dtype=AgiBotWorld_CONFIG["states"][key[state_key_prefix_len:]][
"dtype"
],
)
for key, value in state.items()
},
@@ -77,7 +133,9 @@ def load_local_dataset(
if value.size
else np.zeros(
AgiBotWorld_CONFIG["actions"][key[action_key_prefix_len:]]["shape"],
dtype=AgiBotWorld_CONFIG["actions"][key[action_key_prefix_len:]]["dtype"],
dtype=AgiBotWorld_CONFIG["actions"][key[action_key_prefix_len:]][
"dtype"
],
)
for key, value in action.items()
},
@@ -85,11 +143,5 @@ def load_local_dataset(
for i in range(num_frames)
]
videos = {
f"observation.images.{key}": ob_dir / "videos" / f"{key}_color.mp4"
if "sensor" not in key
else ob_dir / "tactile" / f"{key}.mp4" # HACK: handle tactile videos
for key in AgiBotWorld_CONFIG["images"]
if "depth" not in key
}
videos = get_episode_videos(src_path, task_id, episode_id, AgiBotWorld_CONFIG)
return episode_id, frames, videos