File size: 3,449 Bytes
1bfc3f0 | 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 | """Load ActionCodec2Sparse release packages into a router."""
from __future__ import annotations
import json
from pathlib import Path
from ..backends.actioncodec2.defaults import builtin_action_spaces
from ..backends.actioncodec2.preprocessing import PreprocessingSpec
from ..backends.actioncodec2.presets import format_action_spaces
from ..backends.actioncodec2.release import (
check_gripper_consistency,
check_release_gates,
selected_variants,
)
from ..backends.actioncodec2.sparse import ActionCodec2SparseProcessor, ActionCodec2SparseRouter
TOKENIZER = "actioncodec2_sparse"
def release_layouts(profiles):
"""Built-in layouts whose components are all available in ``profiles``."""
return {
name: preset
for name, preset in builtin_action_spaces().items()
if {component.space for component in preset.components} <= set(profiles)
}
def load_release_router(directory, *, gripper_mode, primitive_order=1, verify_gates, tokenizer=TOKENIZER):
"""Load one gripper-mode/order vocabulary per profile with its cleaning policy."""
manifest, mode, variants = selected_variants(
directory, gripper_mode=gripper_mode, order=primitive_order, tokenizer=tokenizer
)
from ..backends.actioncodec2 import ActionCodec2Processor, ActionCodec2Router
processor_class, router_class = ((ActionCodec2SparseProcessor, ActionCodec2SparseRouter)
if tokenizer == TOKENIZER else (ActionCodec2Processor, ActionCodec2Router))
profiles, preprocessing = {}, {}
for name, variant in variants.items():
entry = (Path(directory) / variant["path"]).resolve()
if not entry.is_relative_to(Path(directory).resolve()):
raise ValueError("release profile path escapes its artifact directory")
processor = processor_class.from_pretrained(entry)
spec = PreprocessingSpec.from_dict(
json.loads((entry / "preprocessing.json").read_text(encoding="utf-8"))
)
check_gripper_consistency(entry, processor, spec, mode)
if verify_gates:
if variant.get("gates_passed") is False:
raise ValueError(f"{entry}: recorded release gates failed")
check_release_gates(entry)
profiles[name], preprocessing[name] = processor, spec
router = router_class.from_processors(
profiles,
token_budgets={name: manifest["codebook_size"] for name in profiles},
presets=list(release_layouts(profiles).values()),
)
metadata = {
"package": manifest.get("package", Path(directory).name),
"tokenizer": tokenizer,
"precision": manifest["precision"],
"codebook_size": manifest["codebook_size"],
"gripper_mode": mode,
"primitive_order": primitive_order,
"variants": list(variants.values()),
}
return router, preprocessing, metadata
def describe_release(directory, *, action_dim, gripper_mode, primitive_order=1, tokenizer=TOKENIZER):
"""Report a release's layouts from its manifest alone."""
_, _, variants = selected_variants(
directory, gripper_mode=gripper_mode, order=primitive_order, tokenizer=tokenizer
)
presets = release_layouts(variants)
return format_action_spaces(
presets, action_dim=action_dim, codec_name="ActionCodec2Sparse" if tokenizer == TOKENIZER else "ActionCodec2", chunk_horizon_decode=tokenizer == TOKENIZER
)
|