File size: 3,414 Bytes
20c251e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | from __future__ import annotations
import subprocess
import sys
from pathlib import Path
from dovla_cil.data.datasets import CILDataset
from dovla_cil.data.lerobot_export import LeRobotExportConfig, export_lerobot_style_dataset
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.utils.io import iter_jsonl, read_json
def _make_toy_cil(tmp_path: Path) -> Path:
dataset_dir = tmp_path / "cil"
generate_cil_dataset(
backend="toy",
tasks=built_in_toy_tasks()[:2],
out_dir=dataset_dir,
num_states_per_task=2,
k=4,
seed=19,
shard_size=8,
inline_observations=True,
)
return dataset_dir
def test_lerobot_style_export_selects_best_record_per_group(tmp_path: Path) -> None:
dataset_dir = _make_toy_cil(tmp_path)
out_dir = tmp_path / "lerobot"
metadata = export_lerobot_style_dataset(
LeRobotExportConfig(
dataset_dir=dataset_dir,
out_dir=out_dir,
max_groups=3,
copy_images=False,
)
)
rows = list(iter_jsonl(out_dir / "train.jsonl"))
dataset = CILDataset(dataset_dir)
assert metadata["schema_version"] == "dovla-cil-lerobot-export/v0"
assert metadata["num_episodes"] == 3
assert len(rows) == 3
assert (out_dir / "tasks.jsonl").exists()
assert read_json(out_dir / "metadata.json") == metadata
for row in rows:
group = dataset.get_group(row["cil"]["group_id"])
assert row["reward"] == max(record.reward.score for record in group)
assert row["cil"]["record_id"] in {record.record_id for record in group}
assert row["task"]
assert row["observation"]["image"] is None
assert "action_chunk" in row
def test_lerobot_style_export_cli_runs_without_network(tmp_path: Path) -> None:
dataset_dir = _make_toy_cil(tmp_path)
out_dir = tmp_path / "cli-export"
result = subprocess.run(
[
sys.executable,
"scripts/export_lerobot_dataset.py",
"--dataset",
str(dataset_dir),
"--out",
str(out_dir),
"--max-groups",
"2",
"--no-images",
],
check=True,
capture_output=True,
text=True,
)
assert "dovla-cil-lerobot-export/v0" in result.stdout
assert len(list(iter_jsonl(out_dir / "train.jsonl"))) == 2
def test_task_balanced_export_covers_tasks_deterministically(tmp_path: Path) -> None:
dataset_dir = _make_toy_cil(tmp_path)
first_out = tmp_path / "balanced-first"
second_out = tmp_path / "balanced-second"
config_kwargs = {
"dataset_dir": dataset_dir,
"max_groups": 4,
"group_sampling": "task_balanced",
"seed": 7,
"copy_images": False,
}
export_lerobot_style_dataset(LeRobotExportConfig(out_dir=first_out, **config_kwargs))
export_lerobot_style_dataset(LeRobotExportConfig(out_dir=second_out, **config_kwargs))
first = list(iter_jsonl(first_out / "train.jsonl"))
second = list(iter_jsonl(second_out / "train.jsonl"))
assert {row["cil"]["task_id"] for row in first} == {
"toy_pick_red_mug",
"toy_put_red_mug_in_blue_bowl",
}
assert [row["cil"]["group_id"] for row in first] == [
row["cil"]["group_id"] for row in second
]
|