| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| from minidreamer.config import load_config |
| from minidreamer.data.replay_buffer import ReplayBuffer |
| from minidreamer.evaluation import evaluate_random_policy, evaluate_world_model |
| from minidreamer.envs.make_env import make_env_from_config |
| from minidreamer.planning.evaluate_planner import evaluate_planner |
| from minidreamer.serialization import load_world_model_checkpoint |
|
|
|
|
| def build_arg_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description="Evaluate MiniDreamer components.") |
| subparsers = parser.add_subparsers(dest="command", required=True) |
|
|
| random_parser = subparsers.add_parser("random", help="Evaluate a random policy.") |
| random_parser.add_argument("--config", type=Path, required=True) |
|
|
| planner_parser = subparsers.add_parser("planner", help="Evaluate a trained planner.") |
| planner_parser.add_argument("--config", type=Path, required=True) |
| planner_parser.add_argument("--checkpoint", type=Path, required=True) |
| planner_parser.add_argument( |
| "--random-action-fraction", |
| type=float, |
| default=0.0, |
| help="Optional evaluation-time action noise. Defaults to 0.0 for a clean planner evaluation.", |
| ) |
|
|
| world_model_parser = subparsers.add_parser("world-model", help="Evaluate held-out world model metrics.") |
| world_model_parser.add_argument("--config", type=Path, required=True) |
| world_model_parser.add_argument("--checkpoint", type=Path, required=True) |
| world_model_parser.add_argument("--replay-dir", type=Path, required=True) |
| world_model_parser.add_argument("--split", type=str, default="val", choices=["train", "val", "test"]) |
| return parser |
|
|
|
|
| def main() -> None: |
| parser = build_arg_parser() |
| args = parser.parse_args() |
| config = load_config(args.config) |
|
|
| if args.command == "random": |
| print(evaluate_random_policy(config)) |
| return |
|
|
| env = make_env_from_config(config, seed=config.get("project", {}).get("seed", 0)) |
| action_dim = env.action_space.n |
| env.close() |
| model, _, metadata = load_world_model_checkpoint(args.checkpoint, action_dim=action_dim, map_location="cpu") |
|
|
| if args.command == "planner": |
| print({ |
| "metadata": metadata, |
| **evaluate_planner( |
| config, |
| model, |
| random_action_fraction=args.random_action_fraction, |
| ), |
| }) |
| return |
|
|
| replay = ReplayBuffer.load(args.replay_dir) |
| print({"metadata": metadata, **evaluate_world_model(config, model, replay, split=args.split)}) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
