| from __future__ import annotations
|
|
|
| import argparse
|
| import copy
|
| import importlib
|
| import json
|
| import sys
|
| from pathlib import Path
|
| from typing import Any
|
|
|
| import yaml
|
|
|
| FINAL_ROOT = Path(__file__).resolve().parent
|
| CODE_ROOTS = [
|
| FINAL_ROOT,
|
| FINAL_ROOT / "baselines",
|
| ]
|
| for root in reversed(CODE_ROOTS):
|
| root_str = str(root)
|
| if root_str not in sys.path:
|
| sys.path.insert(0, root_str)
|
|
|
| ALGORITHM_SPECS = {
|
| "hamworld": {
|
| "label": "HaM-World",
|
| "package": "hamworld",
|
| "default_config": FINAL_ROOT / "hamworld" / "configs" / "compare_dmcontrol.yaml",
|
| "preset_configs": {
|
| "compare_dmcontrol": FINAL_ROOT / "hamworld" / "configs" / "compare_dmcontrol.yaml",
|
| "finger_reacher": FINAL_ROOT / "hamworld" / "configs" / "finger_reacher.yaml",
|
| },
|
| },
|
| "dreamerv3": {
|
| "label": "DreamerV3",
|
| "package": "dreamerv3",
|
| "default_config": FINAL_ROOT / "baselines" / "dreamerv3" / "configs" / "compare_dmcontrol.yaml",
|
| "preset_configs": {
|
| "compare_dmcontrol": FINAL_ROOT / "baselines" / "dreamerv3" / "configs" / "compare_dmcontrol.yaml",
|
| "finger_reacher": FINAL_ROOT / "baselines" / "dreamerv3" / "configs" / "finger_reacher.yaml",
|
| },
|
| },
|
| "tdmpc2": {
|
| "label": "TD-MPC2",
|
| "package": "tdmpc2",
|
| "default_config": FINAL_ROOT / "baselines" / "tdmpc2" / "configs" / "compare_dmcontrol.yaml",
|
| "preset_configs": {
|
| "compare_dmcontrol": FINAL_ROOT / "baselines" / "tdmpc2" / "configs" / "compare_dmcontrol.yaml",
|
| "finger_reacher": FINAL_ROOT / "baselines" / "tdmpc2" / "configs" / "finger_reacher.yaml",
|
| },
|
| },
|
| "ppo": {
|
| "label": "PPO",
|
| "package": "ppo",
|
| "default_config": FINAL_ROOT / "baselines" / "ppo" / "configs" / "low_budget_compare_dmcontrol.yaml",
|
| "preset_configs": {
|
| "compare_dmcontrol": FINAL_ROOT / "baselines" / "ppo" / "configs" / "low_budget_compare_dmcontrol.yaml",
|
| "finger_reacher": FINAL_ROOT / "baselines" / "ppo" / "configs" / "low_budget_finger_reacher.yaml",
|
| },
|
| },
|
| "sac": {
|
| "label": "SAC",
|
| "package": "sac",
|
| "default_config": FINAL_ROOT / "baselines" / "sac" / "configs" / "low_budget_compare_dmcontrol.yaml",
|
| "preset_configs": {
|
| "compare_dmcontrol": FINAL_ROOT / "baselines" / "sac" / "configs" / "low_budget_compare_dmcontrol.yaml",
|
| "finger_reacher": FINAL_ROOT / "baselines" / "sac" / "configs" / "low_budget_finger_reacher.yaml",
|
| },
|
| },
|
| }
|
|
|
| DEFAULT_ALL_ALGORITHMS = ["hamworld", "dreamerv3", "tdmpc2", "ppo", "sac"]
|
|
|
| PRESETS = {
|
| "compare_dmcontrol": ["cartpole_swingup", "cheetah_run"],
|
| "finger_reacher": ["finger_spin", "reacher_easy"],
|
| }
|
|
|
|
|
| def resolve_repo_path(raw_path: str) -> Path:
|
| path = Path(raw_path).expanduser()
|
| if not path.is_absolute():
|
| path = (FINAL_ROOT / path).resolve()
|
| return path
|
|
|
|
|
| def load_config(path: Path) -> dict[str, Any]:
|
| with path.open("r", encoding="utf-8") as handle:
|
| data = yaml.safe_load(handle)
|
| if not isinstance(data, dict):
|
| raise ValueError(f"Config at {path} must decode to a dictionary.")
|
| data.setdefault("_config_path", str(path))
|
| return data
|
|
|
|
|
| def _coerce_value(raw: str) -> Any:
|
| lowered = raw.lower()
|
| if lowered == "true":
|
| return True
|
| if lowered == "false":
|
| return False
|
| if lowered == "none":
|
| return None
|
| try:
|
| return json.loads(raw)
|
| except json.JSONDecodeError:
|
| pass
|
| try:
|
| if "." in raw:
|
| return float(raw)
|
| return int(raw)
|
| except ValueError:
|
| return raw
|
|
|
|
|
| def _set_by_dotted_path(config: dict[str, Any], dotted_key: str, value: Any) -> None:
|
| parts = dotted_key.split(".")
|
| node = config
|
| for part in parts[:-1]:
|
| if part not in node or not isinstance(node[part], dict):
|
| node[part] = {}
|
| node = node[part]
|
| node[parts[-1]] = value
|
|
|
|
|
| def apply_overrides(config: dict[str, Any], overrides: list[str]) -> dict[str, Any]:
|
| updated = copy.deepcopy(config)
|
| for override in overrides:
|
| if "=" not in override:
|
| raise ValueError(f"Invalid override '{override}'. Expected key=value.")
|
| dotted_key, raw_value = override.split("=", 1)
|
| _set_by_dotted_path(updated, dotted_key, _coerce_value(raw_value))
|
| return updated
|
|
|
|
|
| def expand_task_runs(config: dict[str, Any], tasks: list[str]) -> list[dict[str, Any]]:
|
| configured_tasks = config.get("benchmark", {}).get("tasks", [])
|
| task_map = {task.get("id"): task for task in configured_tasks}
|
| missing = [task_id for task_id in tasks if task_id not in task_map]
|
| if missing:
|
| available = ", ".join(sorted(task_map)) if task_map else "<none>"
|
| raise ValueError(f"Task(s) not found in {config['_config_path']}: {', '.join(missing)}. Available: {available}.")
|
|
|
| runs: list[dict[str, Any]] = []
|
| for task_id in tasks:
|
| run_config = copy.deepcopy(config)
|
| run_config["task"] = copy.deepcopy(task_map[task_id])
|
| run_config["experiment"] = copy.deepcopy(config.get("experiment", {}))
|
| base_name = run_config["experiment"].get("name", "experiment")
|
| run_config["experiment"]["resolved_name"] = f"{base_name}_{task_id}"
|
| runs.append(run_config)
|
| return runs
|
|
|
|
|
| def normalize_tasks(tasks: list[str]) -> list[str]:
|
| deduped: list[str] = []
|
| seen: set[str] = set()
|
| for task in tasks:
|
| if task not in seen:
|
| seen.add(task)
|
| deduped.append(task)
|
| return deduped
|
|
|
|
|
| def select_tasks(args: argparse.Namespace) -> list[str]:
|
| if getattr(args, "tasks", None):
|
| return normalize_tasks(args.tasks)
|
| return PRESETS[args.preset]
|
|
|
|
|
| def config_path_for_algorithm(algorithm: str, args: argparse.Namespace) -> Path:
|
| override_key = f"{algorithm}_config".replace("-", "_")
|
| override = getattr(args, override_key, None)
|
| if override:
|
| return resolve_repo_path(override)
|
| preset_name = getattr(args, "preset", None)
|
| if preset_name is not None:
|
| preset_configs = ALGORITHM_SPECS[algorithm].get("preset_configs", {})
|
| if preset_name in preset_configs:
|
| return preset_configs[preset_name]
|
| return ALGORITHM_SPECS[algorithm]["default_config"]
|
|
|
|
|
| def build_overrides(args: argparse.Namespace) -> list[str]:
|
| overrides = list(getattr(args, "override", []) or [])
|
| output_root = resolve_repo_path(args.output_root)
|
| overrides.append(f"experiment.output_dir={output_root}")
|
| if getattr(args, "seed", None) is not None:
|
| overrides.append(f"experiment.seed={args.seed}")
|
| return overrides
|
|
|
|
|
| def describe_task(task: dict[str, Any]) -> str:
|
| if task["suite"] == "dmcontrol":
|
| target = f"{task['domain']}/{task['task']}"
|
| else:
|
| target = task.get("env_name", "<unknown>")
|
| return (
|
| f"{task['id']}: suite={task['suite']}, target={target}, "
|
| f"action_repeat={task.get('action_repeat', 1)}, episode_length={task.get('episode_length', 'n/a')}"
|
| )
|
|
|
|
|
| def list_command(args: argparse.Namespace) -> int:
|
| print("Available presets:")
|
| for preset_name, tasks in PRESETS.items():
|
| print(f" {preset_name}: {', '.join(tasks)}")
|
|
|
| print("\nAlgorithms:")
|
| for algorithm in (args.algo if args.algo != "all" else DEFAULT_ALL_ALGORITHMS):
|
| spec = ALGORITHM_SPECS[algorithm]
|
| config_path = config_path_for_algorithm(algorithm, args)
|
| config = load_config(config_path)
|
| print(f" {spec['label']} ({algorithm})")
|
| print(f" config: {config_path}")
|
| for task in config.get("benchmark", {}).get("tasks", []):
|
| print(f" - {describe_task(task)}")
|
| return 0
|
|
|
|
|
| def build_run_configs(algorithm: str, args: argparse.Namespace, tasks: list[str]) -> tuple[Path, list[dict[str, Any]]]:
|
| if getattr(args, "resume", None):
|
| import torch
|
|
|
| checkpoint_path = resolve_repo_path(args.resume)
|
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
| checkpoint_config = checkpoint["config"]
|
| checkpoint_algorithm = checkpoint_config["experiment"]["algorithm"]
|
| if algorithm != checkpoint_algorithm:
|
| raise ValueError(
|
| f"Resume checkpoint {checkpoint_path} is for algorithm '{checkpoint_algorithm}', not '{algorithm}'."
|
| )
|
| resolved = apply_overrides(checkpoint_config, build_overrides(args))
|
| resolved["_resume_checkpoint"] = str(checkpoint_path)
|
| return checkpoint_path, [resolved]
|
|
|
| config_path = config_path_for_algorithm(algorithm, args)
|
| config = load_config(config_path)
|
| resolved = apply_overrides(config, build_overrides(args))
|
| return config_path, expand_task_runs(resolved, tasks)
|
|
|
|
|
| def print_run_plan(algorithm: str, config_path: Path, runs: list[dict[str, Any]], output_root: Path) -> None:
|
| for run_config in runs:
|
| task = run_config["task"]
|
| seed = run_config["experiment"]["seed"]
|
| print(
|
| f"[plan] algorithm={algorithm} task={task['id']} seed={seed} "
|
| f"config={config_path} outputs={output_root}"
|
| )
|
|
|
|
|
| def train_runs(algorithm: str, runs: list[dict[str, Any]]) -> None:
|
| package_name = ALGORITHM_SPECS[algorithm]["package"]
|
| train_module = importlib.import_module(f"{package_name}.train")
|
| label = ALGORITHM_SPECS[algorithm]["label"]
|
| for run_config in runs:
|
| run_config_copy = copy.deepcopy(run_config)
|
| resume_checkpoint = run_config_copy.pop("_resume_checkpoint", None)
|
| task_id = run_config_copy["task"]["id"]
|
| print(f"[train] {label} -> {task_id}")
|
| train_module.main(config=run_config_copy, resume_checkpoint=resume_checkpoint)
|
|
|
|
|
| def train_command(args: argparse.Namespace) -> int:
|
| if getattr(args, "resume", None):
|
| import torch
|
|
|
| checkpoint_path = resolve_repo_path(args.resume)
|
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
| checkpoint_config = checkpoint["config"]
|
| tasks = [checkpoint_config["task"]["id"]]
|
| algorithms = [checkpoint_config["experiment"]["algorithm"]]
|
| else:
|
| tasks = select_tasks(args)
|
| algorithms = DEFAULT_ALL_ALGORITHMS if args.algo == "all" else [args.algo]
|
| output_root = resolve_repo_path(args.output_root)
|
|
|
| prepared: list[tuple[str, Path, list[dict[str, Any]]]] = []
|
| for algorithm in algorithms:
|
| config_path, runs = build_run_configs(algorithm, args, tasks)
|
| prepared.append((algorithm, config_path, runs))
|
| print_run_plan(algorithm, config_path, runs, output_root)
|
|
|
| if args.dry_run:
|
| return 0
|
|
|
| for algorithm, _, runs in prepared:
|
| train_runs(algorithm, runs)
|
| return 0
|
|
|
|
|
| def add_task_selector_arguments(parser: argparse.ArgumentParser) -> None:
|
| task_group = parser.add_mutually_exclusive_group()
|
| task_group.add_argument(
|
| "--tasks",
|
| nargs="+",
|
| default=None,
|
| help="Task ids to run. Example: --tasks cartpole_swingup cheetah_run",
|
| )
|
| task_group.add_argument(
|
| "--preset",
|
| choices=sorted(PRESETS),
|
| default="compare_dmcontrol",
|
| help="Named task preset. Defaults to compare_dmcontrol.",
|
| )
|
|
|
|
|
| def add_shared_run_arguments(parser: argparse.ArgumentParser, include_algo: bool) -> None:
|
| if include_algo:
|
| parser.add_argument(
|
| "--algo",
|
| choices=["hamworld", "dreamerv3", "tdmpc2", "ppo", "sac", "all"],
|
| default="all",
|
| help="Algorithm selector.",
|
| )
|
| add_task_selector_arguments(parser)
|
| parser.add_argument("--output-root", default="outputs", help="Root directory used for all experiment outputs.")
|
| parser.add_argument("--seed", default=None, type=int, help="Optional seed override for all selected algorithms.")
|
| parser.add_argument("--override", action="append", default=[], help="Extra dotted config override, e.g. training.total_steps=10000")
|
| parser.add_argument("--hamworld-config", dest="hamworld_config", default=None, help="Optional config path override for HaM-World.")
|
| parser.add_argument("--dreamerv3-config", dest="dreamerv3_config", default=None, help="Optional config path override for DreamerV3.")
|
| parser.add_argument("--tdmpc2-config", dest="tdmpc2_config", default=None, help="Optional config path override for TD-MPC2.")
|
| parser.add_argument("--ppo-config", dest="ppo_config", default=None, help="Optional config path override for PPO.")
|
| parser.add_argument("--sac-config", dest="sac_config", default=None, help="Optional config path override for SAC.")
|
| parser.add_argument("--resume", default=None, help="Optional checkpoint path to resume a single run from.")
|
| parser.add_argument("--dry-run", action="store_true", help="Print the resolved run plan without starting training.")
|
|
|
|
|
| def build_parser() -> argparse.ArgumentParser:
|
| parser = argparse.ArgumentParser(description="Unified launcher for the final 5-algorithm paper repository.")
|
| subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
| list_parser = subparsers.add_parser("list", help="Show available presets, tasks, and default configs.")
|
| list_parser.add_argument(
|
| "--algo",
|
| nargs="+",
|
| choices=["hamworld", "dreamerv3", "tdmpc2", "ppo", "sac", "all"],
|
| default=["all"],
|
| help="Algorithm selector.",
|
| )
|
| list_parser.add_argument("--hamworld-config", dest="hamworld_config", default=None, help="Optional config path override for HaM-World.")
|
| list_parser.add_argument("--dreamerv3-config", dest="dreamerv3_config", default=None, help="Optional config path override for DreamerV3.")
|
| list_parser.add_argument("--tdmpc2-config", dest="tdmpc2_config", default=None, help="Optional config path override for TD-MPC2.")
|
| list_parser.add_argument("--ppo-config", dest="ppo_config", default=None, help="Optional config path override for PPO.")
|
| list_parser.add_argument("--sac-config", dest="sac_config", default=None, help="Optional config path override for SAC.")
|
| list_parser.set_defaults(handler=list_command)
|
|
|
| train_parser = subparsers.add_parser("train", help="Train one or more supported algorithms.")
|
| add_shared_run_arguments(train_parser, include_algo=True)
|
| train_parser.set_defaults(handler=train_command)
|
| return parser
|
|
|
|
|
| def main(argv: list[str] | None = None) -> int:
|
| parser = build_parser()
|
| args = parser.parse_args(argv)
|
| try:
|
| if getattr(args, "command", None) == "list" and args.algo == ["all"]:
|
| args.algo = DEFAULT_ALL_ALGORITHMS
|
| return args.handler(args)
|
| except KeyboardInterrupt:
|
| print("Interrupted by user.", file=sys.stderr)
|
| return 130
|
| except Exception as exc:
|
| print(f"Error: {exc}", file=sys.stderr)
|
| return 1
|
|
|
|
|
| if __name__ == "__main__":
|
| raise SystemExit(main())
|
|
|
