lithyeon/franka-pnp-50hz-checkpoints main
imports pathlibPython module-level function call
SHA-2568c33eef2c995c5a02be32bd85350d215c78fc7475e5911eb6bef610893ffa36d
MaleculeH(Db)Md(In)
Evidence
1"""See _CONFIGS for the list of available configs."""
2
3import abc
4from collections.abc import Sequence
5import dataclasses
6import difflib
7import logging
8import pathlib
9from typing import Any, Literal, Protocol, TypeAlias
10
11import etils.epath as epath
12import flax.nnx as nnx
13from typing_extensions import override
14import tyro
15
16import openpi.models.model as _model
17import openpi.models.pi0_config as pi0_config
18import openpi.models.pi0_fast as pi0_fast
19import openpi.models.tokenizer as _tokenizer
20import openpi.policies.aloha_policy as aloha_policy
21import openpi.policies.droid_policy as droid_policy
22import openpi.policies.franka_pnp_policy as franka_pnp_policy
23import openpi.policies.libero_policy as libero_policy
24import openpi.shared.download as _download
25import openpi.shared.normalize as _normalize
26import openpi.training.droid_rlds_dataset as droid_rlds_dataset
27import openpi.training.misc.polaris_config as polaris_config
28import openpi.training.misc.roboarena_config as roboarena_config
29import openpi.training.optimizer as _optimizer
30import openpi.training.weight_loaders as weight_loaders
31import openpi.transforms as _transforms
32
33ModelType: TypeAlias = _model.ModelType
34# Work around a tyro issue with using nnx.filterlib.Filter directly.
35Filter: TypeAlias = nnx.filterlib.Filter
36
37
38@dataclasses.da …
40:13… mines the location of assets (e.g., norm stats) that will be used to set up the data pipeline.
41
42 These assets will be replicated inside the check …
42:55… int under the `assets/asset_id` directory.
43
44 This can be used to load assets from a different checkpoint (e.g., base model checkpoint) or some ot …
1266:2… checkpoint assets intact so prior results remain reproducible.
1267_CONFIGS.append(
1268 dataclasses.replace(
1269 next(config for config in _CONFIGS if config.name == "pi05_franka_pnp_overfit1_full"),
⋯6 lines
1315:22… _name: str) -> TrainConfig:
1316 """Get a config by name."""
1317 if config_name not in _CONFIGS_DICT:
1318 closest = difflib.get_close_matches(config_name, _CONFIGS_DICT.keys(), n=1, cutoff=0.0)
1319 closest_str = f" Did yo …