""" train_curriculum.py =================== """ from __future__ import annotations import argparse import logging import os from dataclasses import dataclass, field from pathlib import Path from typing import Dict, List, Optional import zone_observation as _zo assert _zo.SCHEMA_VERSION == 3, ( f"train_curriculum: zone_observation schema mismatch " f"(expected 3, got {_zo.SCHEMA_VERSION})" ) from zone_observation import ForecastConfig from crop_risk_scorer import RiskWeights # --------------------------------------------------------------------------- # Optional ML imports (graceful degradation) # --------------------------------------------------------------------------- try: import torch _TORCH_AVAILABLE = True except ImportError: _TORCH_AVAILABLE = False try: from weather_forecast_env import make_weather_env from sb3_contrib import MaskablePPO from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.callbacks import BaseCallback _ML_AVAILABLE = True except ImportError as _e: _ML_AVAILABLE = False _ML_IMPORT_ERROR = str(_e) make_weather_env = None MaskablePPO = None Monitor = None BaseCallback = object try: from gru_weather_policy import create_gru_weather_policy_kwargs, get_equivariant_policy_class _GRU_AVAILABLE = True except ImportError: _GRU_AVAILABLE = False create_gru_weather_policy_kwargs = None get_equivariant_policy_class = None try: from physics_dynamics import TemporalDynamicsModel, DynaRolloutBuffer, ZoneStateTensor _DYNAMICS_AVAILABLE = True except ImportError: _DYNAMICS_AVAILABLE = False TemporalDynamicsModel = None DynaRolloutBuffer = None ZoneStateTensor = None # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", handlers=[ logging.FileHandler("training.log"), logging.StreamHandler(), ], ) logger = logging.getLogger(__name__) def set_global_seeds(seed: int) -> None: import random import numpy as np random.seed(seed) np.random.seed(seed) try: import torch torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) except ImportError: pass # --------------------------------------------------------------------------- # Device selection # --------------------------------------------------------------------------- def _select_device(requested: str) -> str: if requested == "cuda": if _TORCH_AVAILABLE and torch.cuda.is_available(): return "cuda" logger.warning("CUDA requested but not available — falling back to CPU.") return "cpu" return requested # --------------------------------------------------------------------------- # Dynamics configuration # --------------------------------------------------------------------------- @dataclass class DynamicsConfig: dynamics_model_path: Optional[str] = None surprise_weight: float = 0.05 update_dynamics_every_n_steps: int = 0 # 0 = frozen fine_tune_epochs: int = 3 transition_buffer_size: int = 10_000 # --------------------------------------------------------------------------- # Curriculum definition # --------------------------------------------------------------------------- def resolve_phase_max_steps(n_zones: int, budget_mode: str, episode_length: int) -> int: n = max(1, int(n_zones)) mode = (budget_mode or "triage").strip().lower() if mode == "legacy": return max(1, int(episode_length)) if mode == "full": return n + 1 if mode == "scarce": return n if mode == "triage": return max(1, n - 1) raise ValueError(f"Unknown budget_mode {budget_mode!r}") @dataclass class CurriculumPhase: name: str total_steps: int episode_length: int n_zones: int risk_weights: RiskWeights budget_mode: str = "triage" learning_rate: float = 3e-4 n_steps: int = 4_096 batch_size: int = 256 n_epochs: int = 10 gamma: float = 0.995 gae_lambda: float = 0.95 clip_range: float = 0.2 ent_coef: float = 0.02 vf_coef: float = 0.5 max_grad_norm: float = 0.5 def resolved_max_steps(self) -> int: return resolve_phase_max_steps(self.n_zones, self.budget_mode, self.episode_length) class WeatherCurriculum: PHASES: Dict[str, CurriculumPhase] = { "normal": CurriculumPhase( name="normal", total_steps=200_000, episode_length=150, n_zones=2, budget_mode="full", risk_weights=RiskWeights(), n_steps=4_096, ent_coef=0.05, ), "monsoon": CurriculumPhase( name="monsoon", total_steps=150_000, episode_length=300, n_zones=3, budget_mode="scarce", risk_weights=RiskWeights( drought_obs_weight=0.40, drought_forecast_weight=0.60, flood_obs_weight=0.70, flood_forecast_weight=0.30, fungi_obs_weight=0.75, fungi_forecast_weight=0.25, supply_drought_weight=0.25, supply_flood_weight=0.50, supply_harvest_pressure_weight=0.25, ), n_steps=4_096, ), "drought": CurriculumPhase( name="drought", total_steps=120_000, episode_length=250, n_zones=3, budget_mode="triage", risk_weights=RiskWeights( drought_obs_weight=0.80, drought_forecast_weight=0.20, flood_obs_weight=0.30, flood_forecast_weight=0.70, fungi_obs_weight=0.55, fungi_forecast_weight=0.45, supply_drought_weight=0.55, supply_flood_weight=0.25, supply_harvest_pressure_weight=0.20, ), n_steps=4_096, ), "heatwave": CurriculumPhase( name="heatwave", total_steps=120_000, episode_length=220, n_zones=4, budget_mode="triage", risk_weights=RiskWeights( drought_obs_weight=0.75, drought_forecast_weight=0.25, flood_obs_weight=0.25, flood_forecast_weight=0.75, fungi_obs_weight=0.50, fungi_forecast_weight=0.50, supply_drought_weight=0.60, supply_flood_weight=0.15, supply_harvest_pressure_weight=0.25, ), n_steps=4_096, ), "humidity": CurriculumPhase( name="humidity", total_steps=100_000, episode_length=200, n_zones=4, budget_mode="triage", risk_weights=RiskWeights( drought_obs_weight=0.30, drought_forecast_weight=0.70, flood_obs_weight=0.50, flood_forecast_weight=0.50, fungi_obs_weight=0.85, fungi_forecast_weight=0.15, supply_drought_weight=0.20, supply_flood_weight=0.30, supply_harvest_pressure_weight=0.50, quality_fungi_weight=0.80, quality_delay_weight=0.20, ), n_steps=4_096, ), } @classmethod def get_phase(cls, name: str) -> CurriculumPhase: if name not in cls.PHASES: raise ValueError( f"Unknown phase '{name}'. Options: {sorted(cls.PHASES)}" ) return cls.PHASES[name] @classmethod def phase_order(cls) -> List[str]: return ["normal", "monsoon", "drought", "heatwave", "humidity"] # --------------------------------------------------------------------------- # Checkpoint callback # --------------------------------------------------------------------------- class CheckpointCallback(BaseCallback): def __init__(self, output_dir: Path, save_freq: int = 25_000) -> None: super().__init__() self.output_dir = output_dir self.save_freq = save_freq self._last_save = 0 def _on_step(self) -> bool: if self.num_timesteps - self._last_save >= self.save_freq: self._last_save = self.num_timesteps path = self.output_dir / f"checkpoint_{self.num_timesteps}.zip" self.model.save(str(path)) logger.info("Checkpoint saved: %s", path.name) return True class RealDataUsageCallback(BaseCallback): """Confirms at runtime -- not just at startup -- what fraction of episodes actually drew real vs synthetic data. info["context_source"] (set by WeatherForecastEnv.reset(), one of "injected"/"real_sampled"/ "synthetic") is only present on the step where an episode boundary triggered an autoreset, so this counts opportunistically rather than on every step. A silent zero real_sampled count with real_data_pkl_path configured is exactly the kind of thing that goes unnoticed for a long time otherwise -- this exists so it can't. """ def __init__(self, log_every: int = 10_000) -> None: super().__init__() self.log_every = log_every self._counts: Dict[str, int] = {} self._last_log = 0 def _on_step(self) -> bool: infos = self.locals.get("infos") if infos: for info in infos: src = info.get("context_source") if isinstance(info, dict) else None if src is not None: self._counts[src] = self._counts.get(src, 0) + 1 if self.num_timesteps - self._last_log >= self.log_every: total = sum(self._counts.values()) if total > 0: logger.info( "RealDataUsageCallback: step=%d episode context_source " "counts so far: %s", self.num_timesteps, dict(self._counts), ) self._last_log = self.num_timesteps return True def _on_training_end(self) -> None: total = sum(self._counts.values()) if total == 0: logger.warning( "RealDataUsageCallback: never observed a context_source in " "any info dict this run -- could not confirm real-vs-synthetic " "mix (this can happen with some VecEnv wrapping; it does not " "necessarily mean sampling failed, but it also cannot confirm " "it succeeded -- verify with a direct env.reset() smoke test " "if this matters for the run)." ) return logger.info( "RealDataUsageCallback: final episode context_source counts: " "%s (%d total, real fraction=%.3f)", dict(self._counts), total, self._counts.get("real_sampled", 0) / total, ) # --------------------------------------------------------------------------- # Dyna callback # --------------------------------------------------------------------------- class DynaCallback(BaseCallback): _OBS_KEYS = ("forecast_precip", "forecast_uncertainty", "zone_belief") def __init__( self, dyna_buffer: "DynaRolloutBuffer", dynamics_cfg: DynamicsConfig, n_zones: int, horizon_days: int, device: str = "cpu", ) -> None: super().__init__() self.dyna_buffer = dyna_buffer self.dynamics_cfg = dynamics_cfg self.n_zones = n_zones self.horizon_days = horizon_days self.device = device self._transition_buffer: list = [] self._tb_max = dynamics_cfg.transition_buffer_size self._bonus_sum = 0.0 self._bonus_count = 0 self._log_freq = 10_000 self._last_log = 0 self._prev_obs: Optional[dict] = None def _obs_to_state_tensor(self, obs: dict) -> Optional["ZoneStateTensor"]: if not all(k in obs for k in self._OBS_KEYS): return None import torch import numpy as np try: precip = np.array(obs["forecast_precip"], dtype=np.float32) uncert = np.array(obs["forecast_uncertainty"], dtype=np.float32) belief = np.array(obs["zone_belief"], dtype=np.float32) if precip.ndim == 2: precip = precip[np.newaxis] # [n_zones, H] -> [1, n_zones, H] if uncert.ndim == 1: uncert = uncert[np.newaxis] # [n_zones] -> [1, n_zones] if belief.ndim == 1: belief = belief[np.newaxis] return ZoneStateTensor( precip=torch.from_numpy(precip).to(self.device), uncertainty=torch.from_numpy(uncert).to(self.device), belief=torch.from_numpy(belief).to(self.device), ) except Exception as e: logger.debug("DynaCallback._obs_to_state_tensor failed: %s", e) return None def _on_step(self) -> bool: try: obs_now = self.locals.get("obs_tensor") or self.locals.get("obs") obs_next = self.locals.get("new_obs") if obs_now is None or obs_next is None: return True # safe: missing locals, skip silently if hasattr(obs_now, "numpy"): if hasattr(obs_now, "items"): obs_now_np = {k: v.cpu().numpy() for k, v in obs_now.items()} else: obs_now_np = {"_raw": obs_now.cpu().numpy()} else: obs_now_np = obs_now if hasattr(obs_next, "items"): obs_next_np = {k: (v.cpu().numpy() if hasattr(v, "cpu") else v) for k, v in obs_next.items()} else: obs_next_np = obs_next curr_state = self._obs_to_state_tensor(obs_now_np) next_state = self._obs_to_state_tensor(obs_next_np) if curr_state is None or next_state is None: return True # safe: obs keys not present yet bonus = self.dyna_buffer.compute_surprise_bonus(curr_state, next_state) bonus_val = float(bonus.item()) bonus_clipped = min(bonus_val, self.dynamics_cfg.surprise_weight) rb = self.model.rollout_buffer if rb is not None and hasattr(rb, "rewards") and rb.rewards is not None: idx = (rb.pos - 1) % rb.buffer_size rb.rewards[idx] += bonus_clipped if self.dynamics_cfg.update_dynamics_every_n_steps > 0: self._transition_buffer.append((curr_state, next_state)) if len(self._transition_buffer) > self._tb_max: self._transition_buffer.pop(0) self._bonus_sum += bonus_clipped self._bonus_count += 1 if self.num_timesteps - self._last_log >= self._log_freq: avg_bonus = ( self._bonus_sum / self._bonus_count if self._bonus_count > 0 else 0.0 ) logger.info( "DynaCallback: step=%d avg_surprise_bonus=%.4f " "buffer_size=%d", self.num_timesteps, avg_bonus, len(self._transition_buffer), ) self._bonus_sum = 0.0 self._bonus_count = 0 self._last_log = self.num_timesteps except Exception as e: logger.debug("DynaCallback._on_step error (non-fatal): %s", e) return True def _on_rollout_end(self) -> None: if ( self.dynamics_cfg.update_dynamics_every_n_steps <= 0 or self.num_timesteps % self.dynamics_cfg.update_dynamics_every_n_steps != 0 or len(self._transition_buffer) < 16 ): return try: from physics_dynamics import DynamicsTrainer dynamics_model = self.dyna_buffer.dynamics import torch import torch.nn.functional as F optimizer = torch.optim.AdamW( dynamics_model.parameters(), lr=1e-4, weight_decay=1e-4 ) dynamics_model.train() pairs = list(self._transition_buffer) batch_size = min(32, len(pairs)) for epoch in range(self.dynamics_cfg.fine_tune_epochs): import random random.shuffle(pairs) total_loss = 0.0 n_batches = 0 # FIX: iterate over all pairs, including the final partial batch for i in range(0, len(pairs), batch_size): batch = pairs[i : i + batch_size] curr_list = [p[0] for p in batch] next_list = [p[1] for p in batch] import torch as _t curr_b = ZoneStateTensor( precip=_t.cat([s.precip for s in curr_list], dim=0), uncertainty=_t.cat([s.uncertainty for s in curr_list], dim=0), belief=_t.cat([s.belief for s in curr_list], dim=0), ) next_b = ZoneStateTensor( precip=_t.cat([s.precip for s in next_list], dim=0), uncertainty=_t.cat([s.uncertainty for s in next_list], dim=0), belief=_t.cat([s.belief for s in next_list], dim=0), ) pred, phys_loss = dynamics_model(curr_b, return_physics_loss=True) data_loss = ( F.mse_loss(pred.precip / 500.0, next_b.precip / 500.0) + F.mse_loss(pred.uncertainty, next_b.uncertainty) + F.mse_loss(pred.belief, next_b.belief) ) loss = data_loss + 0.01 * phys_loss optimizer.zero_grad() loss.backward() _t.nn.utils.clip_grad_norm_(dynamics_model.parameters(), 1.0) optimizer.step() total_loss += loss.item() n_batches += 1 dynamics_model.eval() logger.info( "DynaCallback: fine-tuned dynamics model at step=%d " "avg_loss=%.4f n_transitions=%d", self.num_timesteps, total_loss / max(n_batches, 1), len(self._transition_buffer), ) except Exception as e: logger.warning( "DynaCallback._on_rollout_end fine-tune failed (non-fatal): %s", e ) def _build_dyna_callback( dynamics_cfg: Optional[DynamicsConfig], n_zones: int, horizon_days: int, device: str, ) -> Optional["DynaCallback"]: if dynamics_cfg is None or dynamics_cfg.dynamics_model_path is None: return None if not _DYNAMICS_AVAILABLE: logger.warning( "DynamicsConfig provided but physics_dynamics not installed — " "Dyna augmentation disabled." ) return None model_path = Path(dynamics_cfg.dynamics_model_path) if not model_path.exists(): logger.warning( "Dynamics model not found at %s — Dyna augmentation disabled.", model_path, ) return None try: # FIX: load onto the same device as training to avoid CPU/CUDA mismatch import torch as _torch dynamics_model = TemporalDynamicsModel.load( str(model_path), device=_torch.device(device) ) dynamics_model.eval() dyna_buffer = DynaRolloutBuffer( dynamics=dynamics_model, uncertainty_weight=dynamics_cfg.surprise_weight, ) callback = DynaCallback( dyna_buffer=dyna_buffer, dynamics_cfg=dynamics_cfg, n_zones=n_zones, horizon_days=horizon_days, device=device, ) logger.info( "DynaCallback loaded: model=%s surprise_weight=%.3f " "fine_tune_every=%d", model_path.name, dynamics_cfg.surprise_weight, dynamics_cfg.update_dynamics_every_n_steps, ) return callback except Exception as e: logger.warning( "Failed to build DynaCallback (%s) — Dyna augmentation disabled.", e ) return None # --------------------------------------------------------------------------- # Training # --------------------------------------------------------------------------- def transfer_curriculum_weights( resume_from: str, model: "MaskablePPO", device: str = "auto", ) -> "MaskablePPO": old_model = MaskablePPO.load(resume_from, device=device) old_state = old_model.policy.state_dict() new_state = model.policy.state_dict() transferred, skipped = [], [] merged = {} for key, new_tensor in new_state.items(): old_tensor = old_state.get(key) if old_tensor is not None and old_tensor.shape == new_tensor.shape: merged[key] = old_tensor.clone() transferred.append(key) else: merged[key] = new_tensor skipped.append(key) model.policy.load_state_dict(merged) logger.info( "transfer_curriculum_weights: transferred %d/%d parameter tensors from %s " "(freshly initialized: %s)", len(transferred), len(new_state), resume_from, skipped or "none", ) if not transferred: logger.warning( "transfer_curriculum_weights: transferred ZERO parameters -- the " "architectures are likely genuinely incompatible (e.g. resuming " "from a pre-permutation-invariant checkpoint), not just a normal " "n_zones change. Check resume_from's origin before trusting this run." ) return model def train_phase( phase_name: str, output_dir: Path, resume_from: Optional[str] = None, override_steps: Optional[int] = None, hidden_size: int = 64, device: str = "auto", seed: int = 42, dynamics_cfg: Optional[DynamicsConfig] = None, real_data_pkl: Optional[str] = None, real_data_ratio: Optional[float] = None, real_data_noise_scale: Optional[float] = None, ) -> str: if not _ML_AVAILABLE: raise RuntimeError( f"ML stack not available: {_ML_IMPORT_ERROR}\n" "Install: pip install stable-baselines3 sb3-contrib torch" ) output_dir.mkdir(parents=True, exist_ok=True) models_dir = output_dir / "models" models_dir.mkdir(exist_ok=True) device = _select_device( device if device != "auto" else ("cuda" if _TORCH_AVAILABLE and torch.cuda.is_available() else "cpu") ) phase = WeatherCurriculum.get_phase(phase_name) total_steps = override_steps or phase.total_steps max_steps = phase.resolved_max_steps() full_ceiling = phase.n_zones + 1 logger.info( "Phase=%s steps=%d max_steps=%d (budget_mode=%s, full_ceiling=%d) " "n_zones=%d device=%s must_skip=%s", phase.name, total_steps, max_steps, phase.budget_mode, full_ceiling, phase.n_zones, device, "yes" if max_steps < full_ceiling else "no", ) if max_steps >= full_ceiling and phase.budget_mode not in ("full", "legacy"): logger.warning( "Phase %s: max_steps=%d >= full_ceiling=%d despite budget_mode=%s — " "check resolve_phase_max_steps.", phase.name, max_steps, full_ceiling, phase.budget_mode, ) config_kwargs: Dict[str, object] = dict( n_zones=phase.n_zones, seed=seed, soft_reset=True, max_steps=max_steps, real_data_pkl_path=real_data_pkl, inject_noise=bool(real_data_pkl), ) if real_data_ratio is not None: config_kwargs["real_data_ratio"] = real_data_ratio if real_data_noise_scale is not None: config_kwargs["noise_scale"] = real_data_noise_scale config = ForecastConfig(**config_kwargs) # Train cube v4 TIGGE control: precip_mm length 15. # Holdout persistence rows are still horizon 30. Do NOT copy this # into evaluate_checkpoint_real / backtest_indonesia. if config.real_data_pkl_path: config.horizon_days = 15 phase.risk_weights.attach_to_config(config) if config.real_data_pkl_path: logger.info( "=" * 72 + "\n" "REAL-DATA TRAINING ENABLED for phase=%s\n" " pkl = %s\n" " real_data_ratio = %.2f\n" " inject_noise = %s noise_scale = %.3f\n" " holdout = enforced by RealEpisodeIndex " "(real_episode_sampler.DEFAULT_HOLDOUT_RANGES)\n" + "=" * 72, phase.name, config.real_data_pkl_path, config.real_data_ratio, config.inject_noise, config.noise_scale, ) else: logger.warning( "=" * 72 + "\n" "Phase=%s is 100%% SYNTHETIC -- no --real-data-pkl was provided.\n" "Pass --real-data-pkl to change that.\n" + "=" * 72, phase.name, ) env = Monitor(make_weather_env(config)) # FIX: use the zone-equivariant policy class instead of the generic string if _GRU_AVAILABLE: # See train_kaggle.py's build_model() for the full explanation -- # confirmed empirically this session: features_dim must equal # phase.n_zones * hidden_size * 2 to match GRUWeatherFeaturesExtractor's # actual output shape, or policy construction crashes on the first # forward pass. The factory's own default (128, no n_zones scaling) # only happens to be correct for n_zones == 1. features_dim = phase.n_zones * hidden_size * 2 policy_kwargs = create_gru_weather_policy_kwargs( hidden_size=hidden_size, features_dim=features_dim, ) policy = get_equivariant_policy_class() logger.info( "Using GRU policy (hidden_size=%d, n_zones=%d, features_dim=%d)", hidden_size, phase.n_zones, features_dim, ) else: policy_kwargs = dict(net_arch=dict(pi=[128, 64], vf=[128, 64])) policy = "MultiInputPolicy" logger.info("GRU policy unavailable — using MLP policy (net_arch=128,64)") ppo_kwargs = dict( learning_rate=phase.learning_rate, n_steps=phase.n_steps, batch_size=phase.batch_size, n_epochs=phase.n_epochs, gamma=phase.gamma, gae_lambda=phase.gae_lambda, clip_range=phase.clip_range, ent_coef=phase.ent_coef, vf_coef=phase.vf_coef, max_grad_norm=phase.max_grad_norm, device=device, verbose=1, seed=seed, ) if resume_from: logger.info("Resuming from %s", resume_from) model = MaskablePPO( policy=policy, env=env, policy_kwargs=policy_kwargs, **ppo_kwargs, ) model = transfer_curriculum_weights(resume_from, model, device=device) reset_timesteps = False else: model = MaskablePPO( policy=policy, env=env, policy_kwargs=policy_kwargs, **ppo_kwargs, ) reset_timesteps = True # --- Callbacks --- from stable_baselines3.common.callbacks import CallbackList callbacks = [CheckpointCallback(output_dir), RealDataUsageCallback()] horizon_days = getattr(config, "horizon_days", 14) dyna_cb = _build_dyna_callback( dynamics_cfg=dynamics_cfg, n_zones=phase.n_zones, horizon_days=horizon_days, device=device, ) if dyna_cb is not None: callbacks.append(dyna_cb) logger.info("Dyna augmentation active for phase=%s", phase.name) else: logger.info("Dyna augmentation inactive for phase=%s", phase.name) model.learn( total_timesteps=total_steps, callback=CallbackList(callbacks), reset_num_timesteps=reset_timesteps, use_masking=True, ) final_path = models_dir / f"final_{phase.name}.zip" model.save(str(final_path)) logger.info("Saved final model: %s", final_path) return str(final_path) def train_full_curriculum( output_dir: Path, device: str = "auto", seed: int = 42, dynamics_cfg: Optional[DynamicsConfig] = None, real_data_pkl: Optional[str] = None, real_data_ratio: Optional[float] = None, real_data_noise_scale: Optional[float] = None, ) -> None: """Run all phases in order, chaining each phase from the previous.""" phases = WeatherCurriculum.phase_order() resume = None for phase_name in phases: logger.info("=== Starting phase: %s ===", phase_name) resume = train_phase( phase_name=phase_name, output_dir=output_dir / phase_name, resume_from=resume, device=device, seed=seed, dynamics_cfg=dynamics_cfg, real_data_pkl=real_data_pkl, real_data_ratio=real_data_ratio, real_data_noise_scale=real_data_noise_scale, ) logger.info("=== Completed phase: %s ===", phase_name) # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def main() -> None: p = argparse.ArgumentParser( description="MaskablePPO curriculum trainer for WeatherForecastEnv" ) p.add_argument( "--phase", choices=list(WeatherCurriculum.PHASES) + ["all"], default="normal", help="Curriculum phase to run, or 'all' to run full curriculum.", ) p.add_argument("--output-dir", default="./run", help="Root output directory") p.add_argument("--resume-from", default=None, help="Path to checkpoint .zip") p.add_argument("--steps", type=int, default=None, help="Override total_steps") p.add_argument("--hidden-size", type=int, default=64) p.add_argument("--device", default="auto", help="'cpu', 'cuda', or 'auto'") p.add_argument("--seed", type=int, default=42) p.add_argument( "--dynamics-model", default=None, help="Path to pre-trained TemporalDynamicsModel .pt file. Enables Dyna augmentation.", ) p.add_argument( "--dynamics-weight", type=float, default=0.05, help="Surprise bonus weight per step (only used with --dynamics-model). Default 0.05.", ) p.add_argument( "--dynamics-finetune-every", type=int, default=0, help="Fine-tune dynamics model every N steps. 0=frozen (default).", ) p.add_argument( "--real-data-pkl", default=None, help=( "Path to a historical trajectory cache (see real_episode_sampler.py " "/ RealEpisodeIndex). When set, applies to every phase run this " "invocation (including --phase all): training samples real " "historical episodes with probability --real-data-ratio each " "reset, respecting RealEpisodeIndex's built-in eval-window " "holdout. When NOT set (the default), training is 100%% " "synthetic -- this was true of every checkpoint in this project " "before 2026-08-24; see technical_details.md. Any cache file " "works here regardless of name -- naming convention is " "documentation only, see real_episode_sampler.py." ), ) p.add_argument( "--real-data-ratio", type=float, default=None, help=( "Per-episode probability of sampling real data when " "--real-data-pkl is set. Omit to use " "ForecastConfig.real_data_ratio's own default (0.7)." ), ) p.add_argument( "--real-data-noise-scale", type=float, default=None, help=( "Perturbation magnitude applied to sampled real data. Omit to " "use ForecastConfig.noise_scale's own default (0.05). Noise " "injection is enabled automatically whenever --real-data-pkl " "is set." ), ) args = p.parse_args() set_global_seeds(args.seed) logger.info("Global seeds set to %s (Python / NumPy / PyTorch)", args.seed) output_dir = Path(args.output_dir) dynamics_cfg: Optional[DynamicsConfig] = None if args.dynamics_model is not None: dynamics_cfg = DynamicsConfig( dynamics_model_path=args.dynamics_model, surprise_weight=args.dynamics_weight, update_dynamics_every_n_steps=args.dynamics_finetune_every, ) logger.info( "Dyna config: model=%s weight=%.3f finetune_every=%d", args.dynamics_model, args.dynamics_weight, args.dynamics_finetune_every, ) if args.phase == "all": train_full_curriculum( output_dir=output_dir, device=args.device, seed=args.seed, dynamics_cfg=dynamics_cfg, real_data_pkl=args.real_data_pkl, real_data_ratio=args.real_data_ratio, real_data_noise_scale=args.real_data_noise_scale, ) else: train_phase( phase_name=args.phase, output_dir=output_dir, resume_from=args.resume_from, override_steps=args.steps, hidden_size=args.hidden_size, device=args.device, seed=args.seed, dynamics_cfg=dynamics_cfg, real_data_pkl=args.real_data_pkl, real_data_ratio=args.real_data_ratio, real_data_noise_scale=args.real_data_noise_scale, ) if __name__ == "__main__": main()