Skip to content

data.md

Latest commit

 

History

History
112 lines (86 loc) · 4.37 KB

data.md

File metadata and controls

112 lines (86 loc) · 4.37 KB
Raw

Data & Datasets

The imitation-learning data path runs through the public env, dataset, and policy APIs — collect, expand, train, evaluate:

from ioailab.agents import CuroboPlannerAgent
from ioailab.datasets import DatasetRef, mimic
from ioailab.envs import make_env
from ioailab.agents.policy import OptimizerCfg, Policy

task_id = "GalbotG1-PickCube-v0"
env = make_env(task_id, num_envs=9, headless=True)

# 1. Collect — batch helper exports on env termination/truncation or max_steps.
agent = CuroboPlannerAgent.from_task(task_id)
dataset = env.collect(
    agent=agent, path="data/pick_cube_demos.hdf5", episodes=36, max_steps=1000
)


# Teleop uses an explicit review loop around env.collect(..., episodes=1);
# call dataset.drop() if the just-recorded candidate should be discarded.

# 2. Expand — IsaacLab Mimic, using the task stored on the dataset ref.
dataset = mimic(dataset, episodes=36)

# 3. Train — select one registered policy backend.
policy = Policy.from_backend("robomimic_bc_transformer")
train_cfg = policy.make_train_config(
    output_dir="outputs/pick_cube",
    epochs=20,
    sequence_length=10,
    optimizer=OptimizerCfg(learning_rate=1.0e-4),
)
checkpoint = policy.train(dataset, train_cfg)
inference_checkpoint = policy.export_inference_checkpoint(checkpoint)

# 4. Evaluate — load the inference bundle through the same policy backend.
agent = policy.load_checkpoint(inference_checkpoint)
metrics = env.evaluate(agent=agent, episodes=36)

Built-in backends are robomimic_diffusion, robomimic_bc, robomimic_bc_rnn, and robomimic_bc_transformer. Diffusion export promotes EMA weights; BC export promotes the trained network weights. Both remove optimizer, scheduler, and resume-loop state while preserving the full training checkpoint. The deployment bundle contains model_inference.pth, config.json, and action_norm_params.json under the run's inference/ directory.

Policy selection is not coupled to robomimic. A backend implemented against the Policy contract can be registered before use without editing the central factory:

from ioailab.agents.policy import Policy, register_policy_backend

register_policy_backend("project_policy", ProjectPolicy)
policy = Policy.from_backend("project_policy")

The factory must return a Policy implementing the training, checkpoint-load, and inference-export operations it exposes. Registration is strict: empty, duplicate, or non-Policy factories fail instead of falling back to another model. Run training inside make shell; the container's Isaac Sim Python, PyTorch, robomimic, CUDA, and codec versions are the supported environment.

DatasetRef(path, task_id=...) carries the source task ID as provenance. Pass the registered task ID (e.g. GalbotG1-PickCube-v0) — the dataset helper resolves any Mimic-specific env (GalbotG1-PickCube-Mimic-v0) internally. Tasks own their EnvCfg, MDP terms, recorder config, and any Mimic metadata; there are no script bridges or handwritten dataset fallbacks.

LeRobot v3 export

The dev image installs the optional LeRobot v3 writer (lerobot==0.5.1, --no-deps so it cannot downgrade Isaac Sim's curated numpy/torch/CUDA stack). Verify it inside the container:

python -c "from lerobot.datasets.lerobot_dataset import LeRobotDataset; print(LeRobotDataset)"

Export a staged HDF5 dataset with the explicit LeRobot submodule exporter:

from pathlib import Path

from ioailab.datasets.motion_plan_lerobot import MotionPlanLeRobotExporter

exported = MotionPlanLeRobotExporter(
    hdf5_path=Path("logs/lerobot/stack_cube_motion_plan_staging.hdf5"),
    lerobot_root=Path("logs/lerobot/stack_cube"),
).export()

The staging file sits beside the dataset root (not inside it), and the root must not already exist — LeRobot creates the directory structure itself. Exported features cover action, observation.state, and optional RGB image streams (observation.images.*) as LeRobot video features; depth/RGBD are not exported yet.

YOLO Segmentation Datasets

Warning

REMOTE MACHINE: NOT SUPPORTED. Run the GPU-backed vision baseline only from a local clone on a machine with suitable NVIDIA GPU resources.

ioailab can also generate YOLO-seg datasets directly from task scenes using semantic segmentation cameras. See docs/yolo_seg.md for the full pipeline (dataset generation, label visualization, model training, and inference).