File size: 3,962 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
from pathlib import Path
from dovla_cil.data.datasets import CILDataset, write_cil_collection
from dovla_cil.data.group_sampler import GroupAwareBatchSampler
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.training.collate import collate_cil_records
def _make_toy_dataset(tmp_path: Path) -> CILDataset:
generate_cil_dataset(
backend="toy",
tasks=built_in_toy_tasks()[:3],
out_dir=tmp_path,
num_states_per_task=1,
k=4,
seed=13,
shard_size=8,
inline_observations=True,
)
return CILDataset(tmp_path)
def test_cil_dataset_loads_generated_toy_cil(tmp_path: Path) -> None:
dataset = _make_toy_dataset(tmp_path)
assert len(dataset) == 12
first = dataset[0]
assert dataset.get_group(first.group_id)[0] == first
assert len(list(dataset.iter_groups())) == 3
def test_full_group_sampler_keeps_group_ids_together(tmp_path: Path) -> None:
dataset = _make_toy_dataset(tmp_path)
sampler = GroupAwareBatchSampler(dataset, mode="full_group", batch_groups=1)
batch_indices = next(iter(sampler))
records = [dataset[index] for index in batch_indices]
assert len({record.group_id for record in records}) == 1
assert len(records) == 4
def test_pair_sampler_returns_same_group_reward_ordered_pairs(tmp_path: Path) -> None:
dataset = _make_toy_dataset(tmp_path)
sampler = GroupAwareBatchSampler(
dataset,
mode="pairs",
batch_groups=3,
pair_count_per_group=2,
shuffle=False,
seed=3,
)
batch_indices = next(iter(sampler))
records = [dataset[index] for index in batch_indices]
assert batch_indices.pair_indices
for better_index, worse_index in batch_indices.pair_indices:
better = records[better_index]
worse = records[worse_index]
assert better.group_id == worse.group_id
assert better.reward.score > worse.reward.score
def test_collate_returns_tensors_and_metadata(tmp_path: Path) -> None:
dataset = _make_toy_dataset(tmp_path)
sampler = GroupAwareBatchSampler(dataset, mode="full_group", batch_groups=1)
records = [dataset[index] for index in next(iter(sampler))]
batch = collate_cil_records(records)
assert batch["observations"].shape[0] == len(records)
assert batch["action_features"].shape[0] == len(records)
assert batch["effect_features"].shape[0] == len(records)
assert batch["rewards"].shape == (len(records),)
assert len(batch["instructions"]) == len(records)
assert len(batch["action_chunks"]) == len(records)
assert len(batch["effects"]) == len(records)
assert len(batch["group_ids"]) == len(records)
assert len(batch["candidate_types"]) == len(records)
assert len(batch["failures"]) == len(records)
assert "pair_indices" in batch
def test_zero_copy_collection_loads_disjoint_source_datasets(tmp_path: Path) -> None:
tasks = built_in_toy_tasks()
sources = [tmp_path / "source-a", tmp_path / "source-b"]
for source, task, seed in zip(sources, tasks[:2], (31, 47), strict=False):
generate_cil_dataset(
backend="toy",
tasks=[task],
out_dir=source,
num_states_per_task=2,
k=3,
seed=seed,
shard_size=8,
inline_observations=True,
)
collection_dir = tmp_path / "collection"
write_cil_collection(collection_dir, sources, dataset_name="two-task-collection")
dataset = CILDataset(collection_dir)
assert len(dataset) == 12
assert len(dataset.group_ids) == 4
assert dataset.index.metadata["dataset_name"] == "two-task-collection"
assert dataset.index.metadata["task_count"] == 2
assert {record.metadata["source_dataset"] for record in dataset} == {
str(source.resolve()) for source in sources
}
|