modal_train.py

"""Modal app: BC-warm-started RecurrentPPO RL training on a modest GPU.

Current config: a modest GPU (A10G) + many CPU env workers via SubprocVecEnv.
This workload is env/CPU-bound -- the GPU only does the small LSTM-policy PPO
updates, while the CPUs run the terrarium sim. Benchmarks favored A10G with
64 CPUs and 64 envs over A100 with 32 CPUs.

The smoke flagged that the base torch image LACKS numpy, so it is pinned here.

Usage (run OUTSIDE the sandbox -- these hit the network):
    uv run modal run modal_train.py::smoke      # ~30-60k steps, pipeline check
    uv run modal run modal_train.py::full       # bounded full run (1-2M steps)

Both download the metrics report (and full saves the policy) into
terrarium/artifacts/ locally via the local entrypoint.

Cost discipline: A10G ~= $1.10/hr, billed per second. Pre-calculate
GPU $/hr * wall-clock before any paid run; see the run notes printed locally.
"""

from __future__ import annotations

import json
from pathlib import Path

import modal

APP_NAME = "terrarium-rl-train"
LOCAL_ROOT = Path(__file__).resolve().parent
ARTIFACT_DIR = LOCAL_ROOT / "terrarium" / "artifacts"

# --- image -----------------------------------------------------------------
# debian_slim + pinned RL stack. numpy is explicit (base torch image lacked it
# per the smoke). cloudpickle is needed for SubprocVecEnv spawn workers.
image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "numpy",
        "torch",
        "gymnasium>=1.2.3",
        "stable-baselines3>=2.8.0",
        "sb3-contrib>=2.8.0",
        "cloudpickle",
    )
    # Add the local terrarium package source so rl_env + train_rl import on Modal.
    .add_local_python_source("terrarium")
)

app = modal.App(APP_NAME, image=image)

# Persist trained artifacts on Modal too (belt-and-suspenders; the local
# entrypoint also writes them under terrarium/artifacts/).
volume = modal.Volume.from_name("terrarium-rl-artifacts", create_if_missing=True)
VOL_MOUNT = "/artifacts"


def _train(
    *,
    rl_timesteps: int,
    teacher_episodes: int,
    bc_epochs: int,
    eval_episodes: int,
    max_steps: int,
    n_steps: int,
    n_envs: int,
    seeds: list[int],
    policy_filename: str | None,
    large_world_projection: bool = False,
    map_profiles: list[str] | None = None,
    hidden_size: int = 128,
    lstm_layers: int = 1,
    learning_rate: float = 3e-4,
    ent_coef: float = 0.01,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
    ppo_epochs: int = 4,
    batch_size: int | None = None,
    social_directive_chance: float = 0.0,
    compare_no_command: bool = True,
    skip_bc: bool = False,
) -> dict:
    """Shared body for the GPU functions: run the pipeline, persist on the volume."""
    import time

    import torch

    from terrarium.train_rl import train_and_report

    save_path = None
    if policy_filename:
        save_path = str(Path(VOL_MOUNT) / policy_filename)

    wall_start = time.time()
    report = train_and_report(
        rl_timesteps=rl_timesteps,
        teacher_episodes=teacher_episodes,
        bc_epochs=bc_epochs,
        eval_episodes=eval_episodes,
        max_steps=max_steps,
        n_steps=n_steps,
        n_envs=n_envs,
        seeds=seeds,
        device="auto",
        save_policy_path=save_path,
        hidden_size=hidden_size,
        lstm_layers=lstm_layers,
        learning_rate=learning_rate,
        ent_coef=ent_coef,
        gamma=gamma,
        gae_lambda=gae_lambda,
        ppo_epochs=ppo_epochs,
        batch_size=batch_size,
        large_world_projection=large_world_projection,
        map_profiles=map_profiles,
        social_directive_chance=social_directive_chance,
        compare_no_command=compare_no_command,
        skip_bc=skip_bc,
    )
    report["modal"] = {
        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
        "cuda_available": torch.cuda.is_available(),
        "wall_seconds_remote": round(time.time() - wall_start, 2),
    }
    if save_path:
        volume.commit()
        # ship the policy bytes back so the local entrypoint can write it
        try:
            report["_policy_bytes_b64"] = _read_b64(save_path + ".zip")
            report["_policy_basename"] = policy_filename + ".zip"
        except FileNotFoundError:
            report["_policy_bytes_b64"] = None
    return report


def _read_b64(path: str) -> str:
    import base64

    return base64.b64encode(Path(path).read_bytes()).decode("ascii")


def _parse_csv_ints(raw: str, *, default: list[int]) -> list[int]:
    values = [part.strip() for part in str(raw or "").split(",") if part.strip()]
    return [int(value) for value in values] if values else list(default)


def _parse_csv_strings(raw: str, *, default: list[str]) -> list[str]:
    values = [part.strip() for part in str(raw or "").split(",") if part.strip()]
    return values if values else list(default)


@app.function(gpu="A10G", cpu=8.0, timeout=1200, volumes={VOL_MOUNT: volume})
def smoke_train() -> dict:
    """Cheap Modal smoke: confirm the pipeline runs on GPU under partial obs.

    ~40k steps, 1 seed, 6 parallel env workers. Proves BC warm-start loads, env
    steps on Modal, GPU is used. Derive throughput from wall-clock to size full.
    """
    return _train(
        rl_timesteps=40_000,
        teacher_episodes=24,
        bc_epochs=8,
        eval_episodes=12,
        max_steps=250,
        n_steps=256,
        n_envs=6,
        seeds=[7],
        policy_filename="rl_smoke_policy",
        large_world_projection=True,
        map_profiles=["cover_maze", "risk_reward", "fog_world", "territory_competition", "social_cover"],
        social_directive_chance=0.15,
    )


@app.function(gpu="A10G", cpu=16.0, timeout=7200, volumes={VOL_MOUNT: volume})
def full_train(
    rl_timesteps: int = 250_000,
    seeds: list[int] | None = None,
    candidate: str = "balanced_rich",
) -> dict:
    """Bounded full run: BC warm-start -> RecurrentPPO at GPU scale.

    Defaults are tuned for a two-hour bounded sweep: one seed, rich projected
    maps, and 250k steps/arm. The timeout is the hard ceiling.
    """
    if seeds is None:
        seeds = [7]
    presets = {
        "balanced_rich": {
            "hidden_size": 128,
            "lstm_layers": 1,
            "learning_rate": 3e-4,
            "ent_coef": 0.02,
            "gamma": 0.99,
            "gae_lambda": 0.95,
            "ppo_epochs": 4,
            "map_profiles": ["cover_maze", "risk_reward", "fog_world", "territory_competition", "social_cover"],
        },
        "larger_memory": {
            "hidden_size": 256,
            "lstm_layers": 1,
            "learning_rate": 2e-4,
            "ent_coef": 0.01,
            "gamma": 0.995,
            "gae_lambda": 0.97,
            "ppo_epochs": 4,
            "map_profiles": ["mixed_biomes", "cover_maze", "fog_world", "territory_competition", "social_cover"],
        },
        "survival_1200": {
            "hidden_size": 256,
            "lstm_layers": 1,
            "learning_rate": 2e-4,
            "ent_coef": 0.015,
            "gamma": 0.997,
            "gae_lambda": 0.97,
            "ppo_epochs": 4,
            "map_profiles": ["mixed_biomes", "cover_maze", "risk_reward", "fog_world", "territory_competition", "social_cover", "scarcity", "shelter_maze", "food_vs_heat", "dense_obstacles"],
        },
        "exploration_biased": {
            "hidden_size": 128,
            "lstm_layers": 1,
            "learning_rate": 3e-4,
            "ent_coef": 0.05,
            "gamma": 0.995,
            "gae_lambda": 0.95,
            "ppo_epochs": 5,
            "map_profiles": ["large_open", "curiosity_grove", "cover_maze", "risk_reward", "fog_world", "social_cover"],
        },
    }
    if candidate not in presets:
        raise ValueError(f"unknown candidate {candidate!r}; choose one of {sorted(presets)}")
    config = presets[candidate]
    return _train(
        rl_timesteps=rl_timesteps,
        teacher_episodes=96 if candidate == "survival_1200" else 64,
        bc_epochs=48 if candidate == "survival_1200" else 16,
        eval_episodes=12 if candidate == "survival_1200" else 20,
        max_steps=1200 if candidate == "survival_1200" else 300,
        n_steps=512 if candidate == "survival_1200" else 256,
        n_envs=16 if candidate == "survival_1200" else 8,
        seeds=seeds,
        policy_filename=f"rl_full_policy_{candidate}_{rl_timesteps}",
        large_world_projection=True,
        batch_size=512 if candidate == "survival_1200" else 256,
        social_directive_chance=0.20 if candidate == "survival_1200" else 0.0,
        compare_no_command=False if candidate == "survival_1200" else True,
        **config,
    )


def _pure_reward_train(
    *,
    stages: list[dict],
    eval_episodes: int,
    final_eval_episodes: int,
    final_max_steps: int,
    n_steps: int,
    n_envs: int,
    seed: int,
    policy_filename: str | None,
    hidden_size: int = 256,
    lstm_layers: int = 1,
    learning_rate: float = 2e-4,
    ent_coef: float = 0.02,
    gamma: float = 0.997,
    gae_lambda: float = 0.97,
    ppo_epochs: int = 4,
    batch_size: int | None = None,
    map_profiles: list[str] | None = None,
    social_directive_chance: float = 0.2,
    hierarchical_motor: bool = False,
) -> dict:
    import time

    import torch

    from terrarium.train_rl import train_pure_reward_curriculum_and_report

    save_path = None
    if policy_filename:
        save_path = str(Path(VOL_MOUNT) / policy_filename)
    wall_start = time.time()
    report = train_pure_reward_curriculum_and_report(
        stages=stages,
        eval_episodes=eval_episodes,
        final_eval_episodes=final_eval_episodes,
        final_max_steps=final_max_steps,
        n_steps=n_steps,
        n_envs=n_envs,
        seed=seed,
        device="auto",
        save_policy_path=save_path,
        hidden_size=hidden_size,
        lstm_layers=lstm_layers,
        learning_rate=learning_rate,
        ent_coef=ent_coef,
        gamma=gamma,
        gae_lambda=gae_lambda,
        ppo_epochs=ppo_epochs,
        batch_size=batch_size,
        large_world_projection=True,
        map_profiles=map_profiles,
        social_directive_chance=social_directive_chance,
        hierarchical_motor=hierarchical_motor,
    )
    report["modal"] = {
        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
        "cuda_available": torch.cuda.is_available(),
        "wall_seconds_remote": round(time.time() - wall_start, 2),
    }
    if save_path:
        volume.commit()
        try:
            report["_policy_bytes_b64"] = _read_b64(save_path + ".zip")
            report["_policy_basename"] = policy_filename + ".zip"
        except FileNotFoundError:
            report["_policy_bytes_b64"] = None
    return report


@app.function(gpu="A10G", cpu=16.0, timeout=3600, volumes={VOL_MOUNT: volume})
def pure_reward_stage_train(
    *,
    initial_policy_b64: str | None,
    stage_index: int,
    rl_timesteps: int,
    max_steps: int,
    eval_episodes: int,
    n_steps: int,
    n_envs: int,
    seed: int,
    policy_filename: str,
    map_profiles: list[str],
    social_directive_chance: float,
    hidden_size: int = 256,
    lstm_layers: int = 1,
    learning_rate: float = 2e-4,
    ent_coef: float = 0.025,
    gamma: float = 0.997,
    gae_lambda: float = 0.97,
    ppo_epochs: int = 4,
    batch_size: int | None = 512,
    hierarchical_motor: bool = False,
) -> dict:
    import base64
    import time

    import torch
    from sb3_contrib import RecurrentPPO

    from terrarium.train_rl import build_env, build_model, evaluate_policy_episodes

    device = "cuda" if torch.cuda.is_available() else "cpu"
    vec_env = build_env(
        True,
        max_steps,
        seed + stage_index * 10_000,
        n_envs=n_envs,
        large_world_projection=True,
        map_profiles=map_profiles,
        social_directive_chance=social_directive_chance,
        hierarchical_motor=hierarchical_motor,
    )
    started = time.time()
    initial_loaded = bool(initial_policy_b64)
    if initial_policy_b64:
        initial_path = Path("/tmp") / f"initial_stage_{stage_index}.zip"
        initial_path.write_bytes(base64.b64decode(initial_policy_b64))
        model = RecurrentPPO.load(str(initial_path), env=vec_env, device=device)
    else:
        model = build_model(
            vec_env,
            n_steps=n_steps,
            seed=seed,
            device=device,
            batch_size=batch_size,
            hidden_size=hidden_size,
            lstm_layers=lstm_layers,
            learning_rate=learning_rate,
            ent_coef=ent_coef,
            gamma=gamma,
            gae_lambda=gae_lambda,
            ppo_epochs=ppo_epochs,
        )
    print(
        f"[pure_stage] stage={stage_index} start "
        f"requested_steps={rl_timesteps} max_steps={max_steps} "
        f"n_envs={n_envs} loaded={initial_loaded}",
        flush=True,
    )
    before = int(model.num_timesteps)
    try:
        model.learn(total_timesteps=rl_timesteps, reset_num_timesteps=False, progress_bar=False)
    finally:
        try:
            vec_env.close()
        except Exception:
            pass
    actual = int(model.num_timesteps) - before
    train_wall = time.time() - started
    print(
        f"[pure_stage] stage={stage_index} trained actual_steps={actual} "
        f"wall={train_wall:.1f}s sps={actual / max(1.0, train_wall):.1f}",
        flush=True,
    )
    eval_stats = evaluate_policy_episodes(
        model,
        episodes=eval_episodes,
        max_steps=max_steps,
        condition_on_command=True,
        seed=seed + 777 + stage_index,
        deterministic=False,
        large_world_projection=True,
        map_profiles=map_profiles,
        social_directive_chance=social_directive_chance,
        survival_guard=False,
        hierarchical_motor=hierarchical_motor,
    )
    save_path = str(Path(VOL_MOUNT) / policy_filename)
    model.save(save_path)
    volume.commit()
    policy_bytes = _read_b64(save_path + ".zip")
    return {
        "stage_index": stage_index,
        "requested_rl_timesteps": rl_timesteps,
        "actual_rl_timesteps": actual,
        "max_steps": max_steps,
        "eval_episodes": eval_episodes,
        "map_profiles": map_profiles,
        "social_directive_chance": social_directive_chance,
        "initial_policy_loaded": initial_loaded,
        "hierarchical_motor": bool(hierarchical_motor),
        "train_wall_seconds": round(train_wall, 2),
        "throughput_steps_per_second": round(actual / max(1.0, train_wall), 2),
        "eval": eval_stats,
        "modal": {
            "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
            "cuda_available": torch.cuda.is_available(),
        },
        "_policy_bytes_b64": policy_bytes,
        "_policy_basename": policy_filename + ".zip",
    }


@app.function(gpu="A10G", cpu=8.0, timeout=3600, volumes={VOL_MOUNT: volume})
def pure_reward_final_eval(
    *,
    policy_b64: str,
    final_eval_episodes: int,
    final_max_steps: int,
    seed: int,
    map_profiles: list[str],
    social_directive_chance: float,
    hierarchical_motor: bool = False,
) -> dict:
    import base64
    import time

    import numpy as np
    import torch
    from sb3_contrib import RecurrentPPO

    from terrarium.train_rl import evaluate_command_dataset, evaluate_directive_steering, evaluate_policy_episodes

    started = time.time()
    path = Path("/tmp") / "pure_reward_final_eval_policy.zip"
    path.write_bytes(base64.b64decode(policy_b64))
    model = RecurrentPPO.load(str(path), device="cuda" if torch.cuda.is_available() else "cpu")
    profile_eval: dict[str, object] = {}
    for idx, profile in enumerate(map_profiles):
        profile_eval[profile] = evaluate_policy_episodes(
            model,
            episodes=final_eval_episodes,
            max_steps=final_max_steps,
            condition_on_command=True,
            seed=seed + 20_000 + idx,
            deterministic=False,
            large_world_projection=True,
            map_profiles=[profile],
            social_directive_chance=social_directive_chance,
            survival_guard=False,
            hierarchical_motor=hierarchical_motor,
        )
    directive_eval = evaluate_directive_steering(
        model,
        episodes=max(3, final_eval_episodes),
        max_steps=final_max_steps,
        seed=seed + 30_000,
        large_world_projection=True,
        map_profiles=map_profiles,
        survival_guard=False,
        hierarchical_motor=hierarchical_motor,
    )
    command_dataset_eval = evaluate_command_dataset(
        model,
        limit=320,
        seed=seed + 40_000,
        max_steps=min(final_max_steps, 160),
        large_world_projection=True,
        map_profiles=map_profiles,
        hierarchical_motor=hierarchical_motor,
    )
    values = [v for v in profile_eval.values() if isinstance(v, dict)]
    lengths = [float(v.get("mean_length", 0.0)) for v in values]
    survival_rates = [float(v.get("survival_rate", 0.0)) for v in values]
    summary = {
        "pure_neural_policy_claim": True,
        "reward_only_from_random_weights": True,
        "bc_used": False,
        "teacher_actions_used_for_training": False,
        "survival_guard": False,
        "slow_cognition_conditioned": True,
        "hierarchical_motor": bool(hierarchical_motor),
        "slow_cognition_directives_passed": bool(directive_eval.get("all_cases_pass", False)),
        "command_dataset_size": int(command_dataset_eval.get("dataset_size", 0)),
        "command_dataset_pass_rate": float(command_dataset_eval.get("pass_rate", 0.0)),
        "survives_final_horizon_all_profiles": bool(survival_rates and min(survival_rates) >= 1.0),
        "min_profile_survival_rate": round(min(survival_rates), 4) if survival_rates else 0.0,
        "mean_profile_length": round(float(np.mean(lengths)), 2) if lengths else 0.0,
    }
    return {
        "schema_version": "rl_pure_reward_chunked_final_eval_v1",
        "final_eval_episodes": final_eval_episodes,
        "final_max_steps": final_max_steps,
        "map_profiles": map_profiles,
        "social_directive_chance": social_directive_chance,
        "hierarchical_motor": bool(hierarchical_motor),
        "final_profile_eval": profile_eval,
        "directive_steering_eval": directive_eval,
        "command_dataset_eval": command_dataset_eval,
        "summary": summary,
        "modal": {
            "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
            "cuda_available": torch.cuda.is_available(),
            "wall_seconds_remote": round(time.time() - started, 2),
        },
    }


@app.function(gpu="A10G", cpu=32.0, timeout=43200, volumes={VOL_MOUNT: volume})
def pure_reward_curriculum_train(
    total_scale: str = "mid",
) -> dict:
    profiles = [
        "mixed_biomes", "cover_maze", "risk_reward", "fog_world",
        "territory_competition", "social_cover", "scarcity", "shelter_maze",
        "food_vs_heat", "dense_obstacles",
    ]
    early_profiles = ["mixed_biomes", "cover_maze", "fog_world", "dense_obstacles"]
    if total_scale == "smoke":
        stages = [
            {"rl_timesteps": 2_048, "max_steps": 300, "social_directive_chance": 0.2, "map_profiles": early_profiles},
        ]
        final_eval_episodes = 1
    elif total_scale == "full":
        stages = [
            {"rl_timesteps": 750_000, "max_steps": 300, "social_directive_chance": 0.1, "map_profiles": early_profiles},
            {"rl_timesteps": 1_500_000, "max_steps": 600, "social_directive_chance": 0.2, "map_profiles": profiles},
            {"rl_timesteps": 3_000_000, "max_steps": 1200, "social_directive_chance": 0.25, "map_profiles": profiles},
        ]
        final_eval_episodes = 5
    elif total_scale == "mid":
        stages = [
            {"rl_timesteps": 250_000, "max_steps": 300, "social_directive_chance": 0.1, "map_profiles": early_profiles},
            {"rl_timesteps": 500_000, "max_steps": 600, "social_directive_chance": 0.2, "map_profiles": profiles},
            {"rl_timesteps": 1_000_000, "max_steps": 1200, "social_directive_chance": 0.25, "map_profiles": profiles},
        ]
        final_eval_episodes = 3
    else:
        raise ValueError("total_scale must be one of: smoke, mid, full")
    return _pure_reward_train(
        stages=stages,
        eval_episodes=6 if total_scale != "smoke" else 1,
        final_eval_episodes=final_eval_episodes,
        final_max_steps=1200,
        n_steps=512,
        n_envs=32,
        seed=11,
        policy_filename=f"rl_pure_reward_curriculum_{total_scale}",
        hidden_size=256,
        lstm_layers=1,
        learning_rate=2e-4,
        ent_coef=0.025,
        gamma=0.997,
        gae_lambda=0.97,
        ppo_epochs=4,
        batch_size=1024,
        map_profiles=profiles,
        social_directive_chance=0.25,
    )


@app.function(gpu="A10G", cpu=64.0, timeout=3600, volumes={VOL_MOUNT: volume})
def hybrid_spatial_train(
    *,
    initial_policy_b64: str | None = None,
    rl_timesteps: int = 2_048,
    eval_episodes: int = 2,
    max_steps: int = 160,
    n_steps: int = 128,
    n_envs: int = 4,
    seed: int = 607,
    policy_filename: str = "rl_hybrid_spatial_modal_smoke",
    world_sizes: list[int] | None = None,
    map_profiles: list[str] | None = None,
    radius: int = 8,
    memory_steps: int = 50,
    hidden_size: int = 256,
    lstm_layers: int = 1,
    learning_rate: float = 2.2e-4,
    ent_coef: float = 0.015,
    gamma: float = 0.997,
    gae_lambda: float = 0.97,
    ppo_epochs: int = 4,
    batch_size: int | None = None,
    curriculum: bool = True,
    directive_profile: str = "balanced",
    forbid_follow_directive: bool = False,
    forbid_survive: bool = False,
    directive_reward_scale: float = 1.0,
    social_profile_reward_scale: float = 1.0,
    directive_mismatch_penalty_scale: float = 0.0,
    feature_extractor: str = "combined",
    include_smell: bool = False,
) -> dict:
    """Bounded Modal entrypoint for the hybrid-spatial large-world env.

    This trains from random weights in ``LargeWorldLocalRLEnv`` with dense
    local grid observations plus scalar body/directive features. Defaults are
    intentionally smoke-scale; callers must opt into longer jobs explicitly.
    """
    import base64
    import time

    import numpy as np
    import torch
    from sb3_contrib import RecurrentPPO
    from stable_baselines3.common.callbacks import BaseCallback
    from stable_baselines3.common.utils import get_schedule_fn

    from terrarium.hybrid_eval import evaluate_hybrid_directive_lift
    from terrarium.hybrid_observation import hybrid_observation_shapes
    from terrarium.large_world import LARGE_WORLD_PROFILES
    from terrarium.rl_env import MOTOR_ACTIONS
    from terrarium.train_rl import build_hybrid_spatial_model, build_large_world_local_env

    started = time.time()
    resolved_world_sizes = tuple(int(size) for size in (world_sizes or [256]))
    resolved_profiles = list(map_profiles or ["curiosity_grove", "cover_maze", "fog_world"])
    unknown_profiles = sorted(set(resolved_profiles) - set(LARGE_WORLD_PROFILES))
    if unknown_profiles:
        raise ValueError(f"unknown map_profiles: {unknown_profiles}")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    survival_actions = [
        "survive",
        "follow_directive",
        "seek",
        "seek_food",
        "seek_water",
        "seek_shelter",
        "return_home",
        "rest",
        "rest_near",
        "explore",
    ]
    attention_actions = survival_actions + [
        "inspect",
        "wait_watch",
        "look_left",
        "look_right",
        "look_forward",
        "look_toward",
        "seek_cover",
        "hide",
        "peek",
        "seek_memory",
    ]
    directive_grounding_actions = [
        "survive",
        "seek",
        "seek_food",
        "seek_water",
        "seek_shelter",
        "seek_cover",
        "hide",
        "peek",
        "approach_user",
        "avoid_user",
        "ignore_user",
        "test_from_distance",
        "retreat_from_other",
        "play_near",
        "shadow_other",
        "circle_other",
        "wait_watch",
        "look_toward",
        "inspect",
        "seek_memory",
        "explore",
        "rest",
        "rest_near",
    ]
    if curriculum and rl_timesteps >= 3_000 and directive_profile == "directive_focus":
        phases = [
            ("directive_grounding", max(n_steps * n_envs, int(rl_timesteps * 0.50)), directive_grounding_actions),
            ("full", max(n_steps * n_envs, int(rl_timesteps * 0.50)), None),
        ]
    elif curriculum and rl_timesteps >= 3_000:
        phases = [
            ("survival", max(n_steps * n_envs, int(rl_timesteps * 0.40)), survival_actions),
            ("attention_hide", max(n_steps * n_envs, int(rl_timesteps * 0.30)), attention_actions),
            ("full", max(n_steps * n_envs, int(rl_timesteps * 0.30)), None),
        ]
    else:
        phases = [("full", int(rl_timesteps), None)]
    if forbid_follow_directive:
        phases = [
            (
                name,
                steps,
                tuple(action for action in (actions or MOTOR_ACTIONS) if action != "follow_directive"),
            )
            for name, steps, actions in phases
        ]
    if forbid_survive:
        phases = [
            (
                name,
                steps,
                tuple(action for action in (actions or MOTOR_ACTIONS) if action != "survive"),
            )
            for name, steps, actions in phases
        ]

    def _emit_log(payload: dict) -> None:
        print(json.dumps(payload, sort_keys=True), flush=True)

    class _StageProgressLogger(BaseCallback):
        def __init__(
            self,
            *,
            stage_name: str,
            stage_steps: int,
            stage_started: float,
            interval_steps: int,
            base_timesteps: int,
        ) -> None:
            super().__init__(verbose=0)
            self.stage_name = stage_name
            self.stage_steps = max(1, int(stage_steps))
            self.stage_started = float(stage_started)
            self.interval_steps = max(1, int(interval_steps))
            self.base_timesteps = int(base_timesteps)
            self._last_logged = 0
            self._completion_logged = False

        def _on_step(self) -> bool:
            stage_timesteps = max(0, int(self.num_timesteps) - self.base_timesteps)
            if stage_timesteps >= self.stage_steps and self._completion_logged:
                return True
            if stage_timesteps - self._last_logged < self.interval_steps and stage_timesteps < self.stage_steps:
                return True
            self._last_logged = int(stage_timesteps)
            if stage_timesteps >= self.stage_steps:
                self._completion_logged = True
            elapsed = max(1e-6, time.time() - self.stage_started)
            _emit_log(
                {
                    "event": "hybrid_spatial_train_progress",
                    "stage": self.stage_name,
                    "stage_timesteps": int(stage_timesteps),
                    "stage_requested_timesteps": int(self.stage_steps),
                    "stage_progress": round(min(1.0, stage_timesteps / float(self.stage_steps)), 4),
                    "stage_wall_seconds": round(elapsed, 2),
                    "stage_steps_per_second": round(stage_timesteps / elapsed, 2),
                    "total_model_timesteps": int(model.num_timesteps),
                }
            )
            return True

    def _stage_eval(stage_name: str, stage_index: int) -> dict:
        eval_steps = min(120, int(max_steps))
        eval_env = build_large_world_local_env(
            max_steps=eval_steps,
            seed=seed + 80_000 + stage_index * 1_000,
            n_envs=1,
            world_sizes=resolved_world_sizes,
            map_profiles=resolved_profiles,
            radius=radius,
            memory_steps=memory_steps,
            directive_profile=directive_profile,
            directive_reward_scale=directive_reward_scale,
            social_profile_reward_scale=social_profile_reward_scale,
            directive_mismatch_penalty_scale=directive_mismatch_penalty_scale,
            include_smell=include_smell,
        )
        def _rollout(*, deterministic: bool) -> dict:
            action_counts = np.zeros(len(MOTOR_ACTIONS), dtype=np.int64)
            lengths: list[int] = []
            visited_counts: list[int] = []
            for episode in range(4 if not deterministic else 2):
                obs = eval_env.reset()
                state = None
                episode_start = np.ones((1,), dtype=bool)
                last_info: dict = {}
                for step in range(eval_steps):
                    action, state = model.predict(
                        obs,
                        state=state,
                        episode_start=episode_start,
                        deterministic=deterministic,
                    )
                    action_index = int(action[0])
                    action_counts[action_index] += 1
                    obs, _reward, done, infos = eval_env.step(action)
                    last_info = dict(infos[0]) if infos else {}
                    episode_start = np.asarray(done, dtype=bool)
                    if bool(done[0]):
                        lengths.append(step + 1)
                        break
                else:
                    lengths.append(eval_steps)
                visited_counts.append(int(last_info.get("visited_count", 0)))
            total_actions = int(action_counts.sum())
            top_actions = sorted(
                ((MOTOR_ACTIONS[i], int(action_counts[i])) for i in range(len(MOTOR_ACTIONS)) if int(action_counts[i]) > 0),
                key=lambda item: item[1],
                reverse=True,
            )[:8]
            if total_actions:
                probs = action_counts.astype(np.float64) / float(total_actions)
                entropy = -float(np.sum([p * np.log(p) for p in probs if p > 0.0])) / np.log(len(MOTOR_ACTIONS))
            else:
                entropy = 0.0
            return {
                "deterministic": bool(deterministic),
                "mean_length": round(float(np.mean(lengths)), 2) if lengths else 0.0,
                "mean_visited_cells": round(float(np.mean(visited_counts)), 2) if visited_counts else 0.0,
                "top_actions": dict(top_actions),
                "single_action_share": round(float(top_actions[0][1] / total_actions), 4) if total_actions and top_actions else 0.0,
                "action_entropy_normalized": round(float(entropy), 4),
            }
        try:
            deterministic_report = _rollout(deterministic=True)
            stochastic_report = _rollout(deterministic=False)
        finally:
            try:
                eval_env.close()
            except Exception:
                pass
        report = {
            "event": "hybrid_spatial_stage_eval",
            "stage": stage_name,
            "eval_steps": eval_steps,
            "deterministic": deterministic_report,
            "stochastic": stochastic_report,
        }
        _emit_log(report)
        return report

    _emit_log(
        {
            "event": "hybrid_spatial_train_start",
            "rl_timesteps": int(rl_timesteps),
            "curriculum": bool(curriculum),
            "initial_policy_loaded": bool(initial_policy_b64),
            "phases": [
                {
                    "stage": name,
                    "requested_timesteps": int(steps),
                    "enabled_action_count": len(actions) if actions is not None else len(MOTOR_ACTIONS),
                }
                for name, steps, actions in phases
            ],
            "world_sizes": list(resolved_world_sizes),
            "map_profiles": resolved_profiles,
            "directive_profile": directive_profile,
            "forbid_follow_directive": bool(forbid_follow_directive),
            "forbid_survive": bool(forbid_survive),
            "directive_reward_scale": float(directive_reward_scale),
            "social_profile_reward_scale": float(social_profile_reward_scale),
            "directive_mismatch_penalty_scale": float(directive_mismatch_penalty_scale),
            "feature_extractor": feature_extractor,
            "include_smell": bool(include_smell),
            "n_envs": int(n_envs),
            "n_steps": int(n_steps),
        }
    )

    vec_env = build_large_world_local_env(
        max_steps=max_steps,
        seed=seed,
        n_envs=n_envs,
        world_sizes=resolved_world_sizes,
        map_profiles=resolved_profiles,
        radius=radius,
        memory_steps=memory_steps,
        enabled_motor_actions=phases[0][2],
        directive_profile=directive_profile,
        directive_reward_scale=directive_reward_scale,
        social_profile_reward_scale=social_profile_reward_scale,
        directive_mismatch_penalty_scale=directive_mismatch_penalty_scale,
        include_smell=include_smell,
    )
    initial_loaded = bool(initial_policy_b64)
    if initial_policy_b64:
        initial_path = Path("/tmp") / "hybrid_initial_policy.zip"
        initial_path.write_bytes(base64.b64decode(initial_policy_b64))
        model = RecurrentPPO.load(str(initial_path), env=vec_env, device=device)
        model.learning_rate = learning_rate
        model.lr_schedule = get_schedule_fn(learning_rate)
        model.ent_coef = ent_coef
        model.gamma = gamma
        model.gae_lambda = gae_lambda
        model.n_epochs = ppo_epochs
        if batch_size:
            model.batch_size = int(batch_size)
    else:
        model = build_hybrid_spatial_model(
            vec_env,
            n_steps=n_steps,
            seed=seed,
            device=device,
            batch_size=batch_size,
            hidden_size=hidden_size,
            lstm_layers=lstm_layers,
            learning_rate=learning_rate,
            ent_coef=ent_coef,
            gamma=gamma,
            gae_lambda=gae_lambda,
            ppo_epochs=ppo_epochs,
            feature_extractor=feature_extractor,
        )
    train_started = time.time()
    stage_reports: list[dict] = []
    try:
        for stage_index, (stage_name, stage_steps, enabled_actions) in enumerate(phases):
            if stage_index > 0:
                try:
                    vec_env.close()
                except Exception:
                    pass
                vec_env = build_large_world_local_env(
                    max_steps=max_steps,
                    seed=seed + stage_index * 10_000,
                    n_envs=n_envs,
                    world_sizes=resolved_world_sizes,
                    map_profiles=resolved_profiles,
                    radius=radius,
                    memory_steps=memory_steps,
                    enabled_motor_actions=enabled_actions,
                    directive_profile=directive_profile,
                    directive_reward_scale=directive_reward_scale,
                    social_profile_reward_scale=social_profile_reward_scale,
                    directive_mismatch_penalty_scale=directive_mismatch_penalty_scale,
                    include_smell=include_smell,
                )
                model.set_env(vec_env)
            before_stage = int(model.num_timesteps)
            stage_started = time.time()
            _emit_log(
                {
                    "event": "hybrid_spatial_stage_start",
                    "stage": stage_name,
                    "requested_timesteps": int(stage_steps),
                    "enabled_action_count": len(enabled_actions) if enabled_actions is not None else len(MOTOR_ACTIONS),
                }
            )
            model.learn(
                total_timesteps=int(stage_steps),
                reset_num_timesteps=(stage_index == 0 and not initial_loaded),
                progress_bar=False,
                callback=_StageProgressLogger(
                    stage_name=stage_name,
                    stage_steps=int(stage_steps),
                    stage_started=stage_started,
                    interval_steps=max(n_steps * n_envs * 2, 10_000),
                    base_timesteps=before_stage if (stage_index > 0 or initial_loaded) else 0,
                ),
            )
            stage_eval = _stage_eval(stage_name, stage_index)
            stage_reports.append(
                {
                    "stage": stage_name,
                    "requested_timesteps": int(stage_steps),
                    "actual_timesteps": int(model.num_timesteps) - before_stage,
                    "enabled_action_count": len(enabled_actions) if enabled_actions is not None else len(MOTOR_ACTIONS),
                    "wall_seconds": round(time.time() - stage_started, 2),
                    "deterministic_stage_eval": stage_eval,
                }
            )
    finally:
        try:
            vec_env.close()
        except Exception:
            pass
    train_wall = time.time() - train_started
    actual_steps = int(model.num_timesteps)

    def _final_eval(*, deterministic: bool, seed_offset: int) -> dict:
        eval_env = build_large_world_local_env(
            max_steps=max_steps,
            seed=seed + seed_offset,
            n_envs=1,
            world_sizes=resolved_world_sizes,
            map_profiles=resolved_profiles,
            radius=radius,
            memory_steps=memory_steps,
            directive_profile=directive_profile,
            directive_reward_scale=directive_reward_scale,
            social_profile_reward_scale=social_profile_reward_scale,
            directive_mismatch_penalty_scale=directive_mismatch_penalty_scale,
            include_smell=include_smell,
        )
        returns: list[float] = []
        lengths: list[int] = []
        survived: list[int] = []
        visited_counts: list[int] = []
        known_fractions: list[float] = []
        final_pressures: list[float] = []
        final_health: list[float] = []
        action_counts = np.zeros(len(MOTOR_ACTIONS), dtype=np.int64)
        eval_started = time.time()
        try:
            for _episode in range(int(eval_episodes)):
                obs = eval_env.reset()
                state = None
                episode_start = np.ones((1,), dtype=bool)
                total = 0.0
                steps = 0
                alive = True
                last_info: dict = {}
                for _step in range(int(max_steps)):
                    action, state = model.predict(
                        obs,
                        state=state,
                        episode_start=episode_start,
                        deterministic=deterministic,
                    )
                    action_index = int(action[0])
                    action_counts[action_index] += 1
                    obs, reward, done, infos = eval_env.step(action)
                    last_info = dict(infos[0]) if infos else {}
                    total += float(reward[0])
                    steps += 1
                    episode_start = np.asarray(done, dtype=bool)
                    if bool(done[0]):
                        alive = bool(last_info.get("TimeLimit.truncated", False))
                        break
                returns.append(total)
                lengths.append(steps)
                survived.append(1 if alive and steps >= max_steps else 0)
                visited_counts.append(int(last_info.get("visited_count", 0)))
                known_fractions.append(float(last_info.get("known_world_fraction", 0.0)))
                final_pressures.append(float(last_info.get("survival_pressure", 0.0) or 0.0))
                resources = last_info.get("resources", {})
                final_health.append(float(resources.get("health", 0.0) if isinstance(resources, dict) else 0.0))
        finally:
            try:
                eval_env.close()
            except Exception:
                pass
        eval_wall = time.time() - eval_started
        total_actions = int(action_counts.sum())
        if total_actions:
            probs = action_counts.astype(np.float64) / float(total_actions)
            entropy = -float(np.sum([p * np.log(p) for p in probs if p > 0.0]))
            normalized_entropy = entropy / np.log(len(MOTOR_ACTIONS))
        else:
            normalized_entropy = 0.0
        return {
            "episodes": int(eval_episodes),
            "deterministic": bool(deterministic),
            "mean_return": round(float(np.mean(returns)), 4) if returns else 0.0,
            "mean_length": round(float(np.mean(lengths)), 2) if lengths else 0.0,
            "survival_rate": round(float(np.mean(survived)), 4) if survived else 0.0,
            "mean_visited_cells": round(float(np.mean(visited_counts)), 2) if visited_counts else 0.0,
            "mean_known_world_fraction": round(float(np.mean(known_fractions)), 8) if known_fractions else 0.0,
            "mean_final_survival_pressure": round(float(np.mean(final_pressures)), 4) if final_pressures else 0.0,
            "mean_final_health": round(float(np.mean(final_health)), 4) if final_health else 0.0,
            "action_entropy_normalized": round(float(normalized_entropy), 4),
            "action_distribution": {MOTOR_ACTIONS[i]: int(action_counts[i]) for i in range(len(MOTOR_ACTIONS))},
            "wall_seconds": round(eval_wall, 2),
        }

    stochastic_eval = _final_eval(deterministic=False, seed_offset=50_000)
    deterministic_eval = _final_eval(deterministic=True, seed_offset=60_000)
    neural_steering_eval = evaluate_hybrid_directive_lift(
        model,
        modes=("approach_user", "avoid", "seek_food", "seek_water", "play_near", "hide"),
        repeats=2,
        steps=min(32, int(max_steps)),
        seed=seed + 70_000,
        world_sizes=resolved_world_sizes,
        map_profiles=resolved_profiles,
        radius=radius,
        memory_steps=memory_steps,
        deterministic=False,
        include_smell=include_smell,
    )

    save_path = str(Path(VOL_MOUNT) / policy_filename)
    model.save(save_path)
    volume.commit()
    shapes = hybrid_observation_shapes(radius=radius, include_smell=include_smell)
    report = {
        "schema_version": "rl_hybrid_spatial_modal_report_v1",
        "device": device,
        "config": {
            "rl_timesteps": int(rl_timesteps),
            "actual_rl_timesteps": int(actual_steps),
            "eval_episodes": int(eval_episodes),
            "max_steps": int(max_steps),
            "n_steps": int(n_steps),
            "n_envs": int(n_envs),
            "seed": int(seed),
            "world_sizes": list(resolved_world_sizes),
            "map_profiles": resolved_profiles,
            "radius": int(radius),
            "memory_steps": int(memory_steps),
            "hidden_size": int(hidden_size),
            "lstm_layers": int(lstm_layers),
            "learning_rate": float(learning_rate),
            "ent_coef": float(ent_coef),
            "gamma": float(gamma),
            "gae_lambda": float(gae_lambda),
            "ppo_epochs": int(ppo_epochs),
            "batch_size": batch_size,
            "curriculum": bool(curriculum),
            "initial_policy_loaded": bool(initial_loaded),
            "directive_profile": directive_profile,
            "directive_reward_scale": float(directive_reward_scale),
            "social_profile_reward_scale": float(social_profile_reward_scale),
            "directive_mismatch_penalty_scale": float(directive_mismatch_penalty_scale),
            "feature_extractor": feature_extractor,
            "include_smell": bool(include_smell),
            "forbid_follow_directive": bool(forbid_follow_directive),
            "forbid_survive": bool(forbid_survive),
            "curriculum_stages": stage_reports,
            "observation_shapes": {key: list(value) for key, value in shapes.items()},
        },
        "eval": stochastic_eval,
        "deterministic_eval": deterministic_eval,
        "neural_steering_eval": neural_steering_eval,
        "summary": {
            "hybrid_spatial_observation": True,
            "large_world_local_env": True,
            "fifo_map_memory_steps": int(memory_steps),
            "trained_from_random_weights": not bool(initial_loaded),
            "continued_from_initial_policy": bool(initial_loaded),
            "mean_length": stochastic_eval["mean_length"],
            "survival_rate": stochastic_eval["survival_rate"],
            "mean_visited_cells": stochastic_eval["mean_visited_cells"],
            "deterministic_survival_rate": deterministic_eval["survival_rate"],
            "deterministic_mean_length": deterministic_eval["mean_length"],
            "mean_neural_directive_lift": neural_steering_eval["mean_neural_directive_lift"],
            "neural_directive_pass_rate": neural_steering_eval["neural_pass_rate"],
        },
        "saved_policy": save_path,
        "modal": {
            "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
            "cuda_available": torch.cuda.is_available(),
            "wall_seconds_remote": round(time.time() - started, 2),
            "train_wall_seconds": round(train_wall, 2),
            "eval_wall_seconds": round(float(stochastic_eval["wall_seconds"]) + float(deterministic_eval["wall_seconds"]), 2),
            "throughput_steps_per_second": round(actual_steps / max(1.0, train_wall), 2),
        },
    }
    report["_policy_bytes_b64"] = _read_b64(save_path + ".zip")
    report["_policy_basename"] = policy_filename + ".zip"
    return report


def _hybrid_spatial_benchmark_body(
    *,
    label: str,
    rl_timesteps: int,
    max_steps: int,
    n_steps: int,
    n_envs: int,
    seed: int,
    world_sizes: list[int] | None,
    map_profiles: list[str] | None,
    hidden_size: int,
    batch_size: int | None,
) -> dict:
    import time

    import torch

    from terrarium.large_world import LARGE_WORLD_PROFILES
    from terrarium.train_rl import build_hybrid_spatial_model, build_large_world_local_env

    started = time.time()
    resolved_world_sizes = tuple(int(size) for size in (world_sizes or [256, 512]))
    resolved_profiles = list(map_profiles or ["cover_maze", "social_cover", "curiosity_grove", "fog_world"])
    unknown_profiles = sorted(set(resolved_profiles) - set(LARGE_WORLD_PROFILES))
    if unknown_profiles:
        raise ValueError(f"unknown map_profiles: {unknown_profiles}")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    env = build_large_world_local_env(
        max_steps=max_steps,
        seed=seed,
        n_envs=n_envs,
        world_sizes=resolved_world_sizes,
        map_profiles=resolved_profiles,
        radius=8,
        memory_steps=50,
        directive_profile="social_play",
    )
    model = build_hybrid_spatial_model(
        env,
        n_steps=n_steps,
        seed=seed,
        device=device,
        batch_size=batch_size,
        hidden_size=hidden_size,
        lstm_layers=1,
        learning_rate=8e-5,
        ent_coef=0.02,
        gamma=0.997,
        gae_lambda=0.97,
        ppo_epochs=4,
    )
    train_started = time.time()
    try:
        model.learn(total_timesteps=int(rl_timesteps), progress_bar=False)
    finally:
        env.close()
    train_wall = time.time() - train_started
    actual_steps = int(model.num_timesteps)
    return {
        "schema_version": "hybrid_spatial_gpu_benchmark_v1",
        "label": label,
        "config": {
            "rl_timesteps": int(rl_timesteps),
            "actual_steps": actual_steps,
            "max_steps": int(max_steps),
            "n_steps": int(n_steps),
            "n_envs": int(n_envs),
            "world_sizes": list(resolved_world_sizes),
            "map_profiles": resolved_profiles,
            "hidden_size": int(hidden_size),
            "batch_size": batch_size,
        },
        "modal": {
            "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
            "cuda_available": bool(torch.cuda.is_available()),
            "wall_seconds_remote": round(time.time() - started, 2),
            "train_wall_seconds": round(train_wall, 2),
            "throughput_steps_per_second": round(actual_steps / max(1.0, train_wall), 2),
        },
    }


@app.function(gpu="A10G", cpu=32.0, timeout=1800)
def hybrid_spatial_benchmark_a10g(
    *,
    rl_timesteps: int = 8_192,
    max_steps: int = 400,
    n_steps: int = 128,
    n_envs: int = 32,
    seed: int = 941,
    world_sizes: list[int] | None = None,
    map_profiles: list[str] | None = None,
    hidden_size: int = 256,
    batch_size: int | None = 1024,
) -> dict:
    return _hybrid_spatial_benchmark_body(
        label="A10G",
        rl_timesteps=rl_timesteps,
        max_steps=max_steps,
        n_steps=n_steps,
        n_envs=n_envs,
        seed=seed,
        world_sizes=world_sizes,
        map_profiles=map_profiles,
        hidden_size=hidden_size,
        batch_size=batch_size,
    )


@app.function(gpu="A10G", cpu=16.0, timeout=1800)
def hybrid_spatial_benchmark_a10g_cpu16(
    *,
    rl_timesteps: int = 8_192,
    max_steps: int = 400,
    n_steps: int = 128,
    n_envs: int = 16,
    seed: int = 941,
    world_sizes: list[int] | None = None,
    map_profiles: list[str] | None = None,
    hidden_size: int = 256,
    batch_size: int | None = 512,
) -> dict:
    return _hybrid_spatial_benchmark_body(
        label="A10G_cpu16",
        rl_timesteps=rl_timesteps,
        max_steps=max_steps,
        n_steps=n_steps,
        n_envs=n_envs,
        seed=seed,
        world_sizes=world_sizes,
        map_profiles=map_profiles,
        hidden_size=hidden_size,
        batch_size=batch_size,
    )


@app.function(gpu="A10G", cpu=64.0, timeout=1800)
def hybrid_spatial_benchmark_a10g_cpu64(
    *,
    rl_timesteps: int = 8_192,
    max_steps: int = 400,
    n_steps: int = 128,
    n_envs: int = 64,
    seed: int = 941,
    world_sizes: list[int] | None = None,
    map_profiles: list[str] | None = None,
    hidden_size: int = 256,
    batch_size: int | None = 2048,
) -> dict:
    return _hybrid_spatial_benchmark_body(
        label="A10G_cpu64",
        rl_timesteps=rl_timesteps,
        max_steps=max_steps,
        n_steps=n_steps,
        n_envs=n_envs,
        seed=seed,
        world_sizes=world_sizes,
        map_profiles=map_profiles,
        hidden_size=hidden_size,
        batch_size=batch_size,
    )


@app.function(gpu="A100", cpu=32.0, timeout=1800)
def hybrid_spatial_benchmark_a100(
    *,
    rl_timesteps: int = 8_192,
    max_steps: int = 400,
    n_steps: int = 128,
    n_envs: int = 32,
    seed: int = 941,
    world_sizes: list[int] | None = None,
    map_profiles: list[str] | None = None,
    hidden_size: int = 256,
    batch_size: int | None = 1024,
) -> dict:
    return _hybrid_spatial_benchmark_body(
        label="A100",
        rl_timesteps=rl_timesteps,
        max_steps=max_steps,
        n_steps=n_steps,
        n_envs=n_envs,
        seed=seed,
        world_sizes=world_sizes,
        map_profiles=map_profiles,
        hidden_size=hidden_size,
        batch_size=batch_size,
    )


# --- local entrypoints -----------------------------------------------------
def _persist(report: dict, report_name: str) -> None:
    import base64

    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    policy_b64 = report.pop("_policy_bytes_b64", None)
    policy_name = report.pop("_policy_basename", None)
    if policy_b64 and policy_name:
        (ARTIFACT_DIR / policy_name).write_bytes(base64.b64decode(policy_b64))
        print(f"[modal_train] saved policy -> {ARTIFACT_DIR / policy_name}")
    out = ARTIFACT_DIR / report_name
    out.write_text(json.dumps(report, indent=2, sort_keys=True))
    print(f"[modal_train] wrote {out}")
    print(json.dumps(report.get("summary", {}), indent=2))
    print(f"[modal_train] remote wall: {report.get('modal', {}).get('wall_seconds_remote')}s "
          f"on {report.get('modal', {}).get('gpu')}")


@app.local_entrypoint()
def smoke() -> None:
    report = smoke_train.remote()
    _persist(report, "rl_modal_smoke_report.json")
    wall = report.get("modal", {}).get("wall_seconds_remote") or report.get("wall_seconds", 1.0)
    steps = report["config"]["rl_timesteps"] * 2 * len(report["config"]["seeds"])  # both arms
    sps = steps / max(1.0, wall)
    a10g_per_hr = 1.10
    print("\n=== COST / THROUGHPUT (smoke) ===")
    print(f"smoke wall-clock: {wall:.1f}s  ({steps} agent-steps total -> {sps:.0f} steps/s)")
    print(f"smoke cost: ~${a10g_per_hr * wall / 3600:.3f}")
    for full_steps in (1_000_000, 2_000_000):
        for nseeds in (3,):
            full_total = full_steps * 2 * nseeds
            est_wall = full_total / max(1.0, sps)
            print(f"projected {full_steps:,}-step x {nseeds} seeds: "
                  f"~{est_wall/60:.1f} min -> ~${a10g_per_hr * est_wall / 3600:.2f}")


@app.local_entrypoint()
def full(rl_timesteps: int = 250_000, candidate: str = "balanced_rich") -> None:
    report = full_train.remote(rl_timesteps=rl_timesteps, candidate=candidate)
    _persist(report, f"rl_modal_full_report_{candidate}_{rl_timesteps}.json")
    wall = report.get("modal", {}).get("wall_seconds_remote") or report.get("wall_seconds", 1.0)
    print("\n=== ACTUAL SPEND (full) ===")
    print(f"full wall-clock: {wall:.1f}s -> actual A10G cost ~${1.10 * wall / 3600:.2f}")


@app.local_entrypoint()
def pure_reward(total_scale: str = "mid") -> None:
    report = pure_reward_curriculum_train.remote(total_scale=total_scale)
    _persist(report, f"rl_pure_reward_curriculum_{total_scale}_report.json")
    wall = report.get("modal", {}).get("wall_seconds_remote") or report.get("wall_seconds", 1.0)
    stage_steps = sum(int(stage["rl_timesteps"]) for stage in report["config"]["stages"])
    sps = stage_steps / max(1.0, wall)
    print("\n=== ACTUAL SPEND (pure reward curriculum) ===")
    print(f"scale: {total_scale}")
    print(f"wall-clock: {wall:.1f}s  ({stage_steps:,} agent-steps -> {sps:.0f} steps/s)")
    print(f"A10G cost estimate: ~${1.10 * wall / 3600:.2f}")


@app.local_entrypoint()
def hybrid_benchmark(
    target: str = "a10g32",
    rl_timesteps: int = 8192,
    max_steps: int = 400,
    n_steps: int = 128,
    seed: int = 941,
) -> None:
    targets = {
        "a10g16": hybrid_spatial_benchmark_a10g_cpu16,
        "a10g32": hybrid_spatial_benchmark_a10g,
        "a10g64": hybrid_spatial_benchmark_a10g_cpu64,
        "a100": hybrid_spatial_benchmark_a100,
    }
    if target not in targets:
        raise ValueError(f"target must be one of: {', '.join(sorted(targets))}")
    report = targets[target].remote(
        rl_timesteps=rl_timesteps,
        max_steps=max_steps,
        n_steps=n_steps,
        seed=seed,
    )
    report_name = f"hybrid_spatial_benchmark_{target}_{rl_timesteps}.json"
    _persist(report, report_name)
    modal_report = report.get("modal", {})
    config = report.get("config", {})
    print("\n=== HYBRID SPATIAL BENCHMARK ===")
    print(f"target: {target}")
    print(f"gpu: {modal_report.get('gpu')}")
    print(f"n_envs={config.get('n_envs')} n_steps={config.get('n_steps')} batch_size={config.get('batch_size')}")
    print(f"train wall: {modal_report.get('train_wall_seconds')}s")
    print(f"throughput: {modal_report.get('throughput_steps_per_second')} steps/s")


@app.local_entrypoint()
def hybrid_spatial(
    rl_timesteps: int = 2_048,
    eval_episodes: int = 2,
    max_steps: int = 160,
    n_steps: int = 128,
    n_envs: int = 4,
    world_sizes: str = "256",
    map_profiles: str = "curiosity_grove,cover_maze,fog_world",
    radius: int = 8,
    memory_steps: int = 50,
    seed: int = 607,
    hidden_size: int = 256,
    lstm_layers: int = 1,
    learning_rate: float = 2.2e-4,
    ent_coef: float = 0.015,
    ppo_epochs: int = 4,
    batch_size: int = 0,
    initial_policy_path: str = "",
    phase_mode: str = "curriculum",
    directive_profile: str = "balanced",
    forbid_follow_directive: bool = False,
    forbid_survive: bool = False,
    directive_reward_scale: float = 1.0,
    social_profile_reward_scale: float = 1.0,
    directive_mismatch_penalty_scale: float = 0.0,
    feature_extractor: str = "combined",
    include_smell: bool = False,
    label: str = "smoke",
) -> None:
    """Launch a bounded hybrid-spatial Modal train and persist artifacts locally."""
    import base64

    sizes = _parse_csv_ints(world_sizes, default=[256])
    profiles = _parse_csv_strings(map_profiles, default=["curiosity_grove", "cover_maze", "fog_world"])
    safe_label = "".join(ch if ch.isalnum() or ch in {"_", "-"} else "_" for ch in label).strip("_") or "run"
    policy_name = f"rl_hybrid_spatial_modal_{safe_label}_{rl_timesteps}"
    initial_policy_b64 = None
    if initial_policy_path:
        initial_path = Path(initial_policy_path).expanduser()
        if not initial_path.is_absolute():
            initial_path = (LOCAL_ROOT / initial_path).resolve()
        initial_policy_b64 = base64.b64encode(initial_path.read_bytes()).decode("ascii")
    curriculum = phase_mode != "full_only"
    report = hybrid_spatial_train.remote(
        initial_policy_b64=initial_policy_b64,
        rl_timesteps=rl_timesteps,
        eval_episodes=eval_episodes,
        max_steps=max_steps,
        n_steps=n_steps,
        n_envs=n_envs,
        seed=seed,
        policy_filename=policy_name,
        world_sizes=sizes,
        map_profiles=profiles,
        radius=radius,
        memory_steps=memory_steps,
        hidden_size=hidden_size,
        lstm_layers=lstm_layers,
        learning_rate=learning_rate,
        ent_coef=ent_coef,
        ppo_epochs=ppo_epochs,
        batch_size=batch_size or None,
        curriculum=curriculum,
        directive_profile=directive_profile,
        forbid_follow_directive=forbid_follow_directive,
        forbid_survive=forbid_survive,
        directive_reward_scale=directive_reward_scale,
        social_profile_reward_scale=social_profile_reward_scale,
        directive_mismatch_penalty_scale=directive_mismatch_penalty_scale,
        feature_extractor=feature_extractor,
        include_smell=include_smell,
    )
    _persist(report, f"{policy_name}_report.json")
    wall = report.get("modal", {}).get("wall_seconds_remote") or report.get("wall_seconds", 1.0)
    sps = float(report.get("modal", {}).get("throughput_steps_per_second") or 0.0)
    print("\n=== ACTUAL SPEND (hybrid spatial) ===")
    print(f"label: {safe_label}")
    print(f"world_sizes: {sizes}; profiles: {profiles}")
    print(f"hidden_size={hidden_size}; lstm_layers={lstm_layers}; lr={learning_rate}; ent_coef={ent_coef}; ppo_epochs={ppo_epochs}; batch_size={batch_size or None}; feature_extractor={feature_extractor}; include_smell={include_smell}")
    print(f"initial_policy={initial_policy_path or None}; phase_mode={phase_mode}; curriculum={curriculum}; directive_profile={directive_profile}; forbid_follow_directive={forbid_follow_directive}; forbid_survive={forbid_survive}; directive_reward_scale={directive_reward_scale}; social_profile_reward_scale={social_profile_reward_scale}; directive_mismatch_penalty_scale={directive_mismatch_penalty_scale}")
    print(f"wall-clock: {wall:.1f}s  throughput={sps:.1f} steps/s")
    print(f"A10G cost estimate: ~${1.10 * wall / 3600:.2f}")


@app.local_entrypoint()
def pure_reward_chunked(total_scale: str = "pilot") -> None:
    import base64
    import time

    profiles = [
        "mixed_biomes", "cover_maze", "risk_reward", "fog_world",
        "territory_competition", "social_cover", "scarcity", "shelter_maze",
        "food_vs_heat", "dense_obstacles",
    ]
    early_profiles = ["mixed_biomes", "cover_maze", "fog_world", "dense_obstacles"]
    if total_scale == "pilot":
        chunks = [
            {"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.15, "map_profiles": early_profiles},
            {"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.20, "map_profiles": early_profiles},
        ]
        final_eval_episodes = 1
    elif total_scale == "mid":
        chunks = (
            [{"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.15, "map_profiles": early_profiles} for _ in range(5)]
            + [{"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.20, "map_profiles": profiles} for _ in range(5)]
            + [{"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.25, "map_profiles": profiles} for _ in range(10)]
        )
        final_eval_episodes = 3
    elif total_scale == "full":
        chunks = (
            [{"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.15, "map_profiles": early_profiles} for _ in range(10)]
            + [{"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.20, "map_profiles": profiles} for _ in range(20)]
            + [{"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.25, "map_profiles": profiles} for _ in range(40)]
        )
        final_eval_episodes = 5
    else:
        raise ValueError("total_scale must be one of: pilot, mid, full")

    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    policy_b64: str | None = None
    stage_reports: list[dict] = []
    started = time.time()
    for idx, chunk in enumerate(chunks, start=1):
        policy_name = f"rl_pure_reward_chunked_{total_scale}_stage_{idx:03d}"
        report = pure_reward_stage_train.remote(
            initial_policy_b64=policy_b64,
            stage_index=idx,
            rl_timesteps=int(chunk["rl_timesteps"]),
            max_steps=int(chunk["max_steps"]),
            eval_episodes=2 if total_scale != "pilot" else 1,
            n_steps=512,
            n_envs=16,
            seed=17,
            policy_filename=policy_name,
            map_profiles=list(chunk["map_profiles"]),
            social_directive_chance=float(chunk["social_directive_chance"]),
            hidden_size=256,
            lstm_layers=1,
            learning_rate=2e-4,
            ent_coef=0.025,
            gamma=0.997,
            gae_lambda=0.97,
            ppo_epochs=4,
            batch_size=512,
        )
        policy_b64 = report.pop("_policy_bytes_b64")
        policy_basename = report.pop("_policy_basename")
        (ARTIFACT_DIR / policy_basename).write_bytes(base64.b64decode(policy_b64))
        stage_path = ARTIFACT_DIR / f"{policy_name}_report.json"
        stage_path.write_text(json.dumps(report, indent=2, sort_keys=True))
        stage_reports.append(report)
        print(
            f"[chunked] stage {idx}/{len(chunks)} "
            f"length={report['eval']['mean_length']} "
            f"survival={report['eval']['survival_rate']} "
            f"sps={report['throughput_steps_per_second']} "
            f"saved={ARTIFACT_DIR / policy_basename}",
            flush=True,
        )

    assert policy_b64 is not None
    final = pure_reward_final_eval.remote(
        policy_b64=policy_b64,
        final_eval_episodes=final_eval_episodes,
        final_max_steps=1200,
        seed=17,
        map_profiles=profiles,
        social_directive_chance=0.25,
    )
    combined = {
        "schema_version": "rl_pure_reward_chunked_run_v1",
        "total_scale": total_scale,
        "chunks": chunks,
        "stage_reports": stage_reports,
        "final_eval": final,
        "summary": final["summary"],
        "wall_seconds_local": round(time.time() - started, 2),
    }
    out = ARTIFACT_DIR / f"rl_pure_reward_chunked_{total_scale}_report.json"
    out.write_text(json.dumps(combined, indent=2, sort_keys=True))
    print(f"[chunked] wrote {out}")
    print(json.dumps(final["summary"], indent=2))


@app.local_entrypoint()
def pure_reward_variant(variant: str = "survival_first") -> None:
    import base64
    import time

    profiles = [
        "mixed_biomes", "cover_maze", "risk_reward", "fog_world",
        "territory_competition", "social_cover", "scarcity", "shelter_maze",
        "food_vs_heat", "dense_obstacles",
    ]
    early_profiles = ["mixed_biomes", "cover_maze", "fog_world", "dense_obstacles"]
    variants = {
        "survival_first": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 240, "social_directive_chance": 0.0, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.0, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.10, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.15, "map_profiles": profiles},
            ],
            "ent_coef": 0.008,
            "learning_rate": 2e-4,
            "n_steps": 512,
            "n_envs": 16,
            "final_max_steps": 400,
        },
        "short_curriculum": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 160, "social_directive_chance": 0.10, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 220, "social_directive_chance": 0.15, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.20, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.20, "map_profiles": profiles},
            ],
            "ent_coef": 0.02,
            "learning_rate": 2e-4,
            "n_steps": 256,
            "n_envs": 16,
            "final_max_steps": 400,
        },
        "directive_first": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 240, "social_directive_chance": 0.35, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.35, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.30, "map_profiles": profiles},
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.25, "map_profiles": profiles},
            ],
            "ent_coef": 0.035,
            "learning_rate": 3e-4,
            "n_steps": 512,
            "n_envs": 16,
            "final_max_steps": 400,
        },
        "align_survival": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 240, "social_directive_chance": 0.0, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 300, "social_directive_chance": 0.05, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.10, "map_profiles": profiles},
                {"rl_timesteps": 50_000, "max_steps": 500, "social_directive_chance": 0.15, "map_profiles": profiles},
            ],
            "ent_coef": 0.01,
            "learning_rate": 2e-4,
            "n_steps": 512,
            "n_envs": 16,
            "final_max_steps": 500,
        },
        "align_short": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 140, "social_directive_chance": 0.10, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 220, "social_directive_chance": 0.15, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 320, "social_directive_chance": 0.20, "map_profiles": profiles},
                {"rl_timesteps": 50_000, "max_steps": 500, "social_directive_chance": 0.20, "map_profiles": profiles},
            ],
            "ent_coef": 0.025,
            "learning_rate": 2e-4,
            "n_steps": 256,
            "n_envs": 16,
            "final_max_steps": 500,
        },
        "align_directive": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 240, "social_directive_chance": 0.40, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 320, "social_directive_chance": 0.35, "map_profiles": early_profiles},
                {"rl_timesteps": 50_000, "max_steps": 500, "social_directive_chance": 0.30, "map_profiles": profiles},
                {"rl_timesteps": 50_000, "max_steps": 500, "social_directive_chance": 0.25, "map_profiles": profiles},
            ],
            "ent_coef": 0.04,
            "learning_rate": 3e-4,
            "n_steps": 512,
            "n_envs": 16,
            "final_max_steps": 500,
        },
    }
    if variant not in variants:
        raise ValueError(f"unknown variant {variant!r}; choose one of {sorted(variants)}")
    cfg = variants[variant]
    chunks = list(cfg["chunks"])
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    policy_b64: str | None = None
    stage_reports: list[dict] = []
    started = time.time()
    for idx, chunk in enumerate(chunks, start=1):
        policy_name = f"rl_pure_reward_variant_{variant}_stage_{idx:03d}"
        report = pure_reward_stage_train.remote(
            initial_policy_b64=policy_b64,
            stage_index=idx,
            rl_timesteps=int(chunk["rl_timesteps"]),
            max_steps=int(chunk["max_steps"]),
            eval_episodes=2,
            n_steps=int(cfg["n_steps"]),
            n_envs=int(cfg["n_envs"]),
            seed=101,
            policy_filename=policy_name,
            map_profiles=list(chunk["map_profiles"]),
            social_directive_chance=float(chunk["social_directive_chance"]),
            hidden_size=256,
            lstm_layers=1,
            learning_rate=float(cfg["learning_rate"]),
            ent_coef=float(cfg["ent_coef"]),
            gamma=0.997,
            gae_lambda=0.97,
            ppo_epochs=4,
            batch_size=512,
        )
        policy_b64 = report.pop("_policy_bytes_b64")
        policy_basename = report.pop("_policy_basename")
        (ARTIFACT_DIR / policy_basename).write_bytes(base64.b64decode(policy_b64))
        stage_path = ARTIFACT_DIR / f"{policy_name}_report.json"
        stage_path.write_text(json.dumps(report, indent=2, sort_keys=True))
        stage_reports.append(report)
        print(
            f"[variant:{variant}] stage {idx}/{len(chunks)} "
            f"length={report['eval']['mean_length']} "
            f"unique={report['eval']['mean_unique_cells']} "
            f"survival={report['eval']['survival_rate']} "
            f"entropy={report['eval']['action_entropy_normalized']} "
            f"sps={report['throughput_steps_per_second']}",
            flush=True,
        )

    assert policy_b64 is not None
    final = pure_reward_final_eval.remote(
        policy_b64=policy_b64,
        final_eval_episodes=3,
        final_max_steps=int(cfg["final_max_steps"]),
        seed=101,
        map_profiles=profiles,
        social_directive_chance=0.25,
    )
    combined = {
        "schema_version": "rl_pure_reward_variant_run_v1",
        "variant": variant,
        "config": cfg,
        "stage_reports": stage_reports,
        "final_eval": final,
        "summary": final["summary"],
        "wall_seconds_local": round(time.time() - started, 2),
    }
    out = ARTIFACT_DIR / f"rl_pure_reward_variant_{variant}_report.json"
    out.write_text(json.dumps(combined, indent=2, sort_keys=True))
    print(f"[variant:{variant}] wrote {out}")
    print(json.dumps(final["summary"], indent=2))


@app.local_entrypoint()
def pure_reward_continue(source: str = "align_survival_stage_003", plan: str = "resource_lock") -> None:
    import base64
    import time

    profiles = [
        "mixed_biomes", "cover_maze", "risk_reward", "fog_world",
        "territory_competition", "social_cover", "scarcity", "shelter_maze",
        "food_vs_heat", "dense_obstacles",
    ]
    sources = {
        "align_survival_stage_003": ARTIFACT_DIR / "rl_pure_reward_variant_align_survival_stage_003.zip",
        "align_directive_stage_003": ARTIFACT_DIR / "rl_pure_reward_variant_align_directive_stage_003.zip",
        "align_survival_stage_004": ARTIFACT_DIR / "rl_pure_reward_variant_align_survival_stage_004.zip",
        "resource_lock_stage_003": ARTIFACT_DIR / "rl_pure_reward_continue_align_survival_stage_003_resource_lock_stage_003.zip",
        "directive_lock_stage_003": ARTIFACT_DIR / "rl_pure_reward_continue_align_directive_stage_003_directive_lock_stage_003.zip",
        "explore_lock_stage_003": ARTIFACT_DIR / "rl_pure_reward_continue_align_survival_stage_003_explore_lock_stage_003.zip",
    }
    plans = {
        "resource_lock": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.10},
                {"rl_timesteps": 50_000, "max_steps": 500, "social_directive_chance": 0.15},
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.20},
            ],
            "ent_coef": 0.006,
            "learning_rate": 1.5e-4,
            "n_steps": 512,
            "final_max_steps": 600,
        },
        "directive_lock": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 500, "social_directive_chance": 0.35},
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.30},
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.25},
            ],
            "ent_coef": 0.025,
            "learning_rate": 2e-4,
            "n_steps": 512,
            "final_max_steps": 600,
        },
        "explore_lock": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 500, "social_directive_chance": 0.20},
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.25},
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.25},
            ],
            "ent_coef": 0.06,
            "learning_rate": 2.5e-4,
            "n_steps": 512,
            "final_max_steps": 600,
        },
        "horizon_push": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.20},
                {"rl_timesteps": 50_000, "max_steps": 800, "social_directive_chance": 0.20},
                {"rl_timesteps": 50_000, "max_steps": 1000, "social_directive_chance": 0.25},
            ],
            "ent_coef": 0.008,
            "learning_rate": 1.2e-4,
            "n_steps": 512,
            "final_max_steps": 1000,
        },
        "directive_resource_push": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.30},
                {"rl_timesteps": 50_000, "max_steps": 800, "social_directive_chance": 0.25},
                {"rl_timesteps": 50_000, "max_steps": 1000, "social_directive_chance": 0.25},
            ],
            "ent_coef": 0.018,
            "learning_rate": 1.5e-4,
            "n_steps": 512,
            "final_max_steps": 1000,
        },
        "explore_resource_push": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 600, "social_directive_chance": 0.25},
                {"rl_timesteps": 50_000, "max_steps": 800, "social_directive_chance": 0.25},
                {"rl_timesteps": 50_000, "max_steps": 1000, "social_directive_chance": 0.25},
            ],
            "ent_coef": 0.03,
            "learning_rate": 1.8e-4,
            "n_steps": 512,
            "final_max_steps": 1000,
        },
    }
    if source not in sources:
        raise ValueError(f"unknown source {source!r}; choose one of {sorted(sources)}")
    if plan not in plans:
        raise ValueError(f"unknown plan {plan!r}; choose one of {sorted(plans)}")
    source_path = sources[source]
    if not source_path.exists():
        raise FileNotFoundError(f"missing source policy: {source_path}")
    cfg = plans[plan]
    policy_b64: str | None = base64.b64encode(source_path.read_bytes()).decode("ascii")
    stage_reports: list[dict] = []
    started = time.time()
    for idx, chunk in enumerate(cfg["chunks"], start=1):
        policy_name = f"rl_pure_reward_continue_{source}_{plan}_stage_{idx:03d}"
        report = pure_reward_stage_train.remote(
            initial_policy_b64=policy_b64,
            stage_index=idx,
            rl_timesteps=int(chunk["rl_timesteps"]),
            max_steps=int(chunk["max_steps"]),
            eval_episodes=3,
            n_steps=int(cfg["n_steps"]),
            n_envs=16,
            seed=211,
            policy_filename=policy_name,
            map_profiles=profiles,
            social_directive_chance=float(chunk["social_directive_chance"]),
            hidden_size=256,
            lstm_layers=1,
            learning_rate=float(cfg["learning_rate"]),
            ent_coef=float(cfg["ent_coef"]),
            gamma=0.997,
            gae_lambda=0.97,
            ppo_epochs=4,
            batch_size=512,
        )
        policy_b64 = report.pop("_policy_bytes_b64")
        policy_basename = report.pop("_policy_basename")
        (ARTIFACT_DIR / policy_basename).write_bytes(base64.b64decode(policy_b64))
        stage_path = ARTIFACT_DIR / f"{policy_name}_report.json"
        stage_path.write_text(json.dumps(report, indent=2, sort_keys=True))
        stage_reports.append(report)
        print(
            f"[continue:{source}:{plan}] stage {idx}/{len(cfg['chunks'])} "
            f"length={report['eval']['mean_length']} "
            f"unique={report['eval']['mean_unique_cells']} "
            f"survival={report['eval']['survival_rate']} "
            f"entropy={report['eval']['action_entropy_normalized']} "
            f"sps={report['throughput_steps_per_second']}",
            flush=True,
        )
    assert policy_b64 is not None
    final = pure_reward_final_eval.remote(
        policy_b64=policy_b64,
        final_eval_episodes=5,
        final_max_steps=int(cfg["final_max_steps"]),
        seed=211,
        map_profiles=profiles,
        social_directive_chance=0.25,
    )
    combined = {
        "schema_version": "rl_pure_reward_continue_run_v1",
        "source": source,
        "plan": plan,
        "config": cfg,
        "stage_reports": stage_reports,
        "final_eval": final,
        "summary": final["summary"],
        "wall_seconds_local": round(time.time() - started, 2),
    }
    out = ARTIFACT_DIR / f"rl_pure_reward_continue_{source}_{plan}_report.json"
    out.write_text(json.dumps(combined, indent=2, sort_keys=True))
    print(f"[continue:{source}:{plan}] wrote {out}")
    print(json.dumps(final["summary"], indent=2))


@app.local_entrypoint()
def hierarchical_motor(variant: str = "balanced") -> None:
    import base64
    import time

    profiles = [
        "mixed_biomes", "cover_maze", "risk_reward", "fog_world",
        "territory_competition", "social_cover", "scarcity", "shelter_maze",
        "food_vs_heat", "dense_obstacles",
    ]
    configs = {
        "balanced": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.20},
                {"rl_timesteps": 50_000, "max_steps": 800, "social_directive_chance": 0.25},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.30},
            ],
            "ent_coef": 0.02,
            "learning_rate": 2e-4,
        },
        "survival": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.05},
                {"rl_timesteps": 50_000, "max_steps": 800, "social_directive_chance": 0.10},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.15},
            ],
            "ent_coef": 0.006,
            "learning_rate": 1.5e-4,
        },
        "social": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 400, "social_directive_chance": 0.45},
                {"rl_timesteps": 50_000, "max_steps": 800, "social_directive_chance": 0.40},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.35},
            ],
            "ent_coef": 0.03,
            "learning_rate": 2.5e-4,
        },
        "expanded": {
            "chunks": [
                {"rl_timesteps": 60_000, "max_steps": 400, "social_directive_chance": 0.55},
                {"rl_timesteps": 60_000, "max_steps": 800, "social_directive_chance": 0.50},
                {"rl_timesteps": 80_000, "max_steps": 1200, "social_directive_chance": 0.45},
            ],
            "ent_coef": 0.028,
            "learning_rate": 2.2e-4,
        },
        "expanded_v2": {
            "chunks": [
                {"rl_timesteps": 75_000, "max_steps": 400, "social_directive_chance": 0.60},
                {"rl_timesteps": 75_000, "max_steps": 800, "social_directive_chance": 0.55},
                {"rl_timesteps": 100_000, "max_steps": 1200, "social_directive_chance": 0.50},
            ],
            "ent_coef": 0.030,
            "learning_rate": 2.2e-4,
        },
    }
    if variant not in configs:
        raise ValueError(f"unknown variant {variant!r}; choose one of {sorted(configs)}")
    cfg = configs[variant]
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    policy_b64: str | None = None
    stage_reports: list[dict] = []
    started = time.time()
    for idx, chunk in enumerate(cfg["chunks"], start=1):
        policy_name = f"rl_hierarchical_motor_{variant}_stage_{idx:03d}"
        report = pure_reward_stage_train.remote(
            initial_policy_b64=policy_b64,
            stage_index=idx,
            rl_timesteps=int(chunk["rl_timesteps"]),
            max_steps=int(chunk["max_steps"]),
            eval_episodes=4,
            n_steps=512,
            n_envs=16,
            seed=307,
            policy_filename=policy_name,
            map_profiles=profiles,
            social_directive_chance=float(chunk["social_directive_chance"]),
            hidden_size=256,
            lstm_layers=1,
            learning_rate=float(cfg["learning_rate"]),
            ent_coef=float(cfg["ent_coef"]),
            gamma=0.997,
            gae_lambda=0.97,
            ppo_epochs=4,
            batch_size=512,
            hierarchical_motor=True,
        )
        policy_b64 = report.pop("_policy_bytes_b64")
        policy_basename = report.pop("_policy_basename")
        (ARTIFACT_DIR / policy_basename).write_bytes(base64.b64decode(policy_b64))
        stage_path = ARTIFACT_DIR / f"{policy_name}_report.json"
        stage_path.write_text(json.dumps(report, indent=2, sort_keys=True))
        stage_reports.append(report)
        print(
            f"[hierarchical:{variant}] stage {idx}/{len(cfg['chunks'])} "
            f"length={report['eval']['mean_length']} "
            f"survival={report['eval']['survival_rate']} "
            f"unique={report['eval']['mean_unique_cells']} "
            f"commands_motor={report['eval']['action_distribution']} "
            f"sps={report['throughput_steps_per_second']}",
            flush=True,
        )
    assert policy_b64 is not None
    final = pure_reward_final_eval.remote(
        policy_b64=policy_b64,
        final_eval_episodes=5,
        final_max_steps=1200,
        seed=307,
        map_profiles=profiles,
        social_directive_chance=0.30,
        hierarchical_motor=True,
    )
    combined = {
        "schema_version": "rl_hierarchical_motor_run_v1",
        "variant": variant,
        "config": cfg,
        "stage_reports": stage_reports,
        "final_eval": final,
        "summary": final["summary"],
        "wall_seconds_local": round(time.time() - started, 2),
    }
    out = ARTIFACT_DIR / f"rl_hierarchical_motor_{variant}_report.json"
    out.write_text(json.dumps(combined, indent=2, sort_keys=True))
    print(f"[hierarchical:{variant}] wrote {out}")
    print(json.dumps(final["summary"], indent=2))


@app.local_entrypoint()
def hierarchical_motor_continue(source: str = "survival", plan: str = "consolidate") -> None:
    import base64
    import time

    profiles = [
        "mixed_biomes", "cover_maze", "risk_reward", "fog_world",
        "territory_competition", "social_cover", "scarcity", "shelter_maze",
        "food_vs_heat", "dense_obstacles",
    ]
    sources = {
        "survival": ARTIFACT_DIR / "rl_hierarchical_motor_survival_stage_003.zip",
        "social": ARTIFACT_DIR / "rl_hierarchical_motor_social_stage_003.zip",
        "balanced": ARTIFACT_DIR / "rl_hierarchical_motor_balanced_stage_003.zip",
        "expanded": ARTIFACT_DIR / "rl_hierarchical_motor_expanded_stage_003.zip",
        "survival_consolidated": ARTIFACT_DIR / "rl_hierarchical_motor_continue_survival_consolidate_stage_003.zip",
        "social_consolidated": ARTIFACT_DIR / "rl_hierarchical_motor_continue_social_social_consolidate_stage_003.zip",
    }
    plans = {
        "consolidate": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.20},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.20},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.25},
            ],
            "ent_coef": 0.004,
            "learning_rate": 1.0e-4,
        },
        "social_consolidate": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.35},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.30},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.30},
            ],
            "ent_coef": 0.012,
            "learning_rate": 1.2e-4,
        },
        "interesting_probe": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.35},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.40},
            ],
            "ent_coef": 0.018,
            "learning_rate": 8.0e-5,
        },
        "stable_social_probe": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.30},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.35},
            ],
            "ent_coef": 0.010,
            "learning_rate": 8.0e-5,
        },
        "expanded_consolidate": {
            "chunks": [
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.35},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.35},
                {"rl_timesteps": 50_000, "max_steps": 1200, "social_directive_chance": 0.30},
            ],
            "ent_coef": 0.008,
            "learning_rate": 8.0e-5,
        },
    }
    if source not in sources:
        raise ValueError(f"unknown source {source!r}; choose one of {sorted(sources)}")
    if plan not in plans:
        raise ValueError(f"unknown plan {plan!r}; choose one of {sorted(plans)}")
    source_path = sources[source]
    if not source_path.exists():
        raise FileNotFoundError(f"missing source policy: {source_path}")
    cfg = plans[plan]
    policy_b64: str | None = base64.b64encode(source_path.read_bytes()).decode("ascii")
    stage_reports: list[dict] = []
    started = time.time()
    for idx, chunk in enumerate(cfg["chunks"], start=1):
        policy_name = f"rl_hierarchical_motor_continue_{source}_{plan}_stage_{idx:03d}"
        report = pure_reward_stage_train.remote(
            initial_policy_b64=policy_b64,
            stage_index=idx,
            rl_timesteps=int(chunk["rl_timesteps"]),
            max_steps=1200,
            eval_episodes=5,
            n_steps=512,
            n_envs=16,
            seed=409,
            policy_filename=policy_name,
            map_profiles=profiles,
            social_directive_chance=float(chunk["social_directive_chance"]),
            hidden_size=256,
            lstm_layers=1,
            learning_rate=float(cfg["learning_rate"]),
            ent_coef=float(cfg["ent_coef"]),
            gamma=0.997,
            gae_lambda=0.97,
            ppo_epochs=4,
            batch_size=512,
            hierarchical_motor=True,
        )
        policy_b64 = report.pop("_policy_bytes_b64")
        policy_basename = report.pop("_policy_basename")
        (ARTIFACT_DIR / policy_basename).write_bytes(base64.b64decode(policy_b64))
        stage_path = ARTIFACT_DIR / f"{policy_name}_report.json"
        stage_path.write_text(json.dumps(report, indent=2, sort_keys=True))
        stage_reports.append(report)
        print(
            f"[hierarchical_continue:{source}:{plan}] stage {idx}/{len(cfg['chunks'])} "
            f"length={report['eval']['mean_length']} "
            f"survival={report['eval']['survival_rate']} "
            f"unique={report['eval']['mean_unique_cells']} "
            f"commands_motor={report['eval']['action_distribution']} "
            f"sps={report['throughput_steps_per_second']}",
            flush=True,
        )
    assert policy_b64 is not None
    final = pure_reward_final_eval.remote(
        policy_b64=policy_b64,
        final_eval_episodes=8,
        final_max_steps=1200,
        seed=409,
        map_profiles=profiles,
        social_directive_chance=0.25,
        hierarchical_motor=True,
    )
    combined = {
        "schema_version": "rl_hierarchical_motor_continue_run_v1",
        "source": source,
        "plan": plan,
        "config": cfg,
        "stage_reports": stage_reports,
        "final_eval": final,
        "summary": final["summary"],
        "wall_seconds_local": round(time.time() - started, 2),
    }
    out = ARTIFACT_DIR / f"rl_hierarchical_motor_continue_{source}_{plan}_report.json"
    out.write_text(json.dumps(combined, indent=2, sort_keys=True))
    print(f"[hierarchical_continue:{source}:{plan}] wrote {out}")
    print(json.dumps(final["summary"], indent=2))