from __future__ import annotations import argparse from pathlib import Path from typing import Any import numpy as np import torch from torch.nn.utils import clip_grad_norm_ from tqdm import trange from minidreamer.config import ensure_run_dirs, load_config, merge_dicts, save_config from minidreamer.data.collect_random import collect_bootstrap_dataset 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.models.world_model import WorldModel from minidreamer.planning.cem import DiscreteCEMPlanner from minidreamer.planning.evaluate_planner import evaluate_planner from minidreamer.serialization import save_world_model_checkpoint from minidreamer.utils.common import get_device, seed_everything, write_json, write_jsonl def train_world_model_updates( model: WorldModel, replay: ReplayBuffer, optimizer: torch.optim.Optimizer, config: dict[str, Any], num_updates: int, device: torch.device, ) -> list[dict[str, float]]: if num_updates <= 0: return [] model.train() logs: list[dict[str, float]] = [] progress = trange(num_updates, desc="world-model-updates", leave=False) for _ in progress: batch = ReplayBuffer.batch_to_torch(replay.sample_sequences(split="train"), device=device) losses = model.compute_losses(batch, config) optimizer.zero_grad(set_to_none=True) losses["loss"].backward() clip_grad_norm_(model.parameters(), float(config["training"].get("grad_clip_norm", 100.0))) optimizer.step() log_row = { "loss": float(losses["loss"].detach().cpu()), "reward_loss": float(losses["reward_loss"].cpu()), "done_loss": float(losses["done_loss"].cpu()), "kl_loss": float(losses["kl_loss"].cpu()), "recon_loss": float(losses["recon_loss"].cpu()), } logs.append(log_row) progress.set_postfix({key: f"{value:.3f}" for key, value in log_row.items()}) return logs def optimizer_to_device(optimizer: torch.optim.Optimizer, device: torch.device) -> None: for state in optimizer.state.values(): for key, value in state.items(): if torch.is_tensor(value): state[key] = value.to(device) def load_training_state( checkpoint_path: str | Path, config: dict[str, Any], action_dim: int, device: torch.device, ) -> tuple[dict[str, Any], WorldModel, torch.optim.Optimizer, dict[str, Any]]: payload = torch.load(checkpoint_path, map_location=device, weights_only=False) resolved_config = merge_dicts(payload["config"], config) model = WorldModel.from_config(resolved_config, action_dim=action_dim).to(device) model.load_state_dict(payload["model_state"]) optimizer = torch.optim.Adam(model.parameters(), lr=float(resolved_config["training"]["lr"])) optimizer_state = payload.get("optimizer_state") if optimizer_state is not None: optimizer.load_state_dict(optimizer_state) optimizer_to_device(optimizer, device) return resolved_config, model, optimizer, payload.get("metadata", {}) def find_existing_run_artifacts(base_dir: str | Path) -> list[Path]: base = Path(base_dir) if not base.exists(): return [] artifact_files = [ base / "metrics" / "run_summary.json", base / "metrics" / "train_metrics.jsonl", base / "metrics" / "eval_metrics.jsonl", base / "checkpoints" / "world_model_latest.pt", base / "replay" / "metadata.json", ] found = [path for path in artifact_files if path.exists()] if found: return found for subdir_name in ("checkpoints", "metrics", "replay"): subdir = base / subdir_name if subdir.exists(): for child in subdir.iterdir(): found.append(child) break return found def collect_planner_steps( env, replay: ReplayBuffer, model: WorldModel, planner: DiscreteCEMPlanner, num_steps: int, random_action_fraction: float, rng: np.random.Generator, ) -> dict[str, int]: collected_steps = 0 episodes = 0 success_episodes = 0 model.eval() while collected_steps < num_steps: obs, _ = env.reset() observations = [obs] actions: list[int] = [] rewards: list[float] = [] terminated_flags: list[float] = [] truncated_flags: list[float] = [] done_flags: list[float] = [] terminated = False truncated = False with torch.no_grad(): state = model.posterior_step(model.initial_state(1), None, obs, sample=False) while not (terminated or truncated): if rng.random() < random_action_fraction: action = int(env.action_space.sample()) else: action = planner.plan(state).action obs, reward, terminated, truncated, _ = env.step(action) actions.append(action) rewards.append(float(reward)) terminated_flags.append(float(terminated)) truncated_flags.append(float(truncated)) done_flags.append(float(terminated or truncated)) observations.append(obs) collected_steps += 1 if terminated or truncated: break state = model.posterior_step(state, action, obs, sample=False) replay.add_episode( obs=np.asarray(observations, dtype=np.float32), actions=np.asarray(actions, dtype=np.int64), rewards=np.asarray(rewards, dtype=np.float32), terminated=np.asarray(terminated_flags, dtype=np.float32), truncated=np.asarray(truncated_flags, dtype=np.float32), done=np.asarray(done_flags, dtype=np.float32), ) episodes += 1 success_episodes += int(bool(terminated and np.sum(rewards) > 0.0)) return { "env_steps": collected_steps, "episodes": episodes, "success_episodes": success_episodes, } def run_training( config: dict[str, Any], output_dir: str | Path, replay_dir: str | Path | None = None, resume_checkpoint: str | Path | None = None, allow_overwrite_existing_output: bool = False, ) -> dict[str, Any]: seed = config.get("project", {}).get("seed", 0) seed_everything(seed) existing_artifacts = find_existing_run_artifacts(output_dir) if existing_artifacts and resume_checkpoint is None and not allow_overwrite_existing_output: preview = ", ".join(str(path) for path in existing_artifacts[:3]) raise FileExistsError( f"Refusing to overwrite existing run directory '{output_dir}'. " f"Found existing artifacts: {preview}. " "Choose a new --output-dir, resume with --resume-checkpoint, " "or pass --allow-overwrite-existing-output to overwrite intentionally." ) run_dirs = ensure_run_dirs(output_dir) device = get_device(config.get("training", {}).get("device")) env = make_env_from_config(config, seed=seed) action_dim = env.action_space.n env.close() if replay_dir is not None and Path(replay_dir).exists(): replay = ReplayBuffer.load(replay_dir) collection_summary = {"replay_loaded": replay.summary()} else: replay, collection_summary = collect_bootstrap_dataset(config, output_dir=run_dirs["replay"], seed=seed) resume_metadata: dict[str, Any] = {} if resume_checkpoint is not None: config, model, optimizer, resume_metadata = load_training_state( checkpoint_path=resume_checkpoint, config=config, action_dim=action_dim, device=device, ) else: model = WorldModel.from_config(config, action_dim=action_dim).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=float(config["training"]["lr"])) save_config(config, run_dirs["base"] / "resolved_config.yaml") training_logs: list[dict[str, float]] = [] evaluation_logs: list[dict[str, float]] = [] train_collect_ratio = float(config["collection"].get("train_collect_ratio", 1.0)) total_updates_budget = int(config["training"]["train_steps"]) if resume_checkpoint is not None: updates_done = int(resume_metadata.get("updates_done", 0)) checkpoint_env_steps = int(resume_metadata.get("env_steps", 0)) if replay.env_steps > checkpoint_env_steps and updates_done < total_updates_budget: collect_steps_per_iteration = max(1, int(config["collection"].get("collect_steps_per_iteration", 1))) per_iteration_updates = int( config["collection"].get( "gradient_updates_per_iteration", round(collect_steps_per_iteration * train_collect_ratio), ) ) missed_iterations = max(0, round((replay.env_steps - checkpoint_env_steps) / collect_steps_per_iteration)) catch_up_updates = min(total_updates_budget - updates_done, per_iteration_updates * missed_iterations) catch_up_logs = train_world_model_updates(model, replay, optimizer, config, catch_up_updates, device) training_logs.extend(catch_up_logs) updates_done += len(catch_up_logs) else: initial_updates = min(total_updates_budget, max(1, int(round(replay.env_steps * train_collect_ratio)))) training_logs.extend(train_world_model_updates(model, replay, optimizer, config, initial_updates, device)) updates_done = len(training_logs) comparison_budgets = config.get("comparison", {}).get("env_steps", [replay.env_steps]) target_env_steps = int(max(comparison_budgets)) rng = np.random.default_rng(seed) env = make_env_from_config(config, seed=seed) planner = DiscreteCEMPlanner.from_config(model, env.action_space.n, config) eval_every_steps = int(config["evaluation"].get("eval_every_env_steps", target_env_steps)) next_eval_step = replay.env_steps while replay.env_steps < target_env_steps and updates_done < total_updates_budget: collect_steps = min( int(config["collection"]["collect_steps_per_iteration"]), target_env_steps - replay.env_steps, ) collection_row = collect_planner_steps( env, replay, model, planner, num_steps=collect_steps, random_action_fraction=float(config["collection"].get("random_action_fraction_after_planner", 0.0)), rng=rng, ) updates = int(config["collection"].get("gradient_updates_per_iteration", round(collection_row["env_steps"] * train_collect_ratio))) updates = min(updates, total_updates_budget - updates_done) training_logs.extend(train_world_model_updates(model, replay, optimizer, config, updates, device)) updates_done = len(training_logs) replay.save(run_dirs["replay"]) if replay.env_steps >= next_eval_step: world_model_metrics = evaluate_world_model(config, model, replay, split="val", max_episodes=10) planner_metrics = evaluate_planner(config, model, episodes=min(10, config["evaluation"]["episodes"]), seed=seed) random_metrics = evaluate_random_policy(config, episodes=min(10, config["evaluation"]["episodes"]), seed=seed) eval_row = { "env_steps": replay.env_steps, "updates_done": updates_done, **{f"world_model/{key}": value for key, value in world_model_metrics.items()}, **{f"planner/{key}": value for key, value in planner_metrics.items()}, **{f"random/{key}": value for key, value in random_metrics.items()}, } evaluation_logs.append(eval_row) next_eval_step += eval_every_steps save_world_model_checkpoint( run_dirs["checkpoints"] / f"world_model_env_steps_{replay.env_steps}.pt", model, config, optimizer=optimizer, metadata={"env_steps": replay.env_steps, "updates_done": updates_done}, ) env.close() save_world_model_checkpoint( run_dirs["checkpoints"] / "world_model_latest.pt", model, config, optimizer=optimizer, metadata={"env_steps": replay.env_steps, "updates_done": updates_done}, ) write_json(run_dirs["metrics"] / "collection_summary.json", collection_summary) write_jsonl(run_dirs["metrics"] / "train_metrics.jsonl", training_logs) write_jsonl(run_dirs["metrics"] / "eval_metrics.jsonl", evaluation_logs) summary = { "replay": replay.summary(), "updates_done": updates_done, "device": str(device), } write_json(run_dirs["metrics"] / "run_summary.json", summary) return summary def build_arg_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Train the MiniDreamer world model.") parser.add_argument("--config", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--replay-dir", type=Path, default=None, help="Optional existing replay directory.") parser.add_argument("--resume-checkpoint", type=Path, default=None, help="Optional checkpoint to resume from.") parser.add_argument( "--allow-overwrite-existing-output", action="store_true", help="Allow overwriting an existing run directory when not resuming.", ) return parser def main() -> None: parser = build_arg_parser() args = parser.parse_args() config = load_config(args.config) summary = run_training( config, args.output_dir, replay_dir=args.replay_dir, resume_checkpoint=args.resume_checkpoint, allow_overwrite_existing_output=args.allow_overwrite_existing_output, ) print(summary) if __name__ == "__main__": main()