experiments_dynamic.py

"""Dynamic obstacles: the experiment where a world model finally earns its keep.

On the DynamicGridWorld (a vertical wall bar sliding through P=4 phases, moving
every step) we compare the project's three Dyna-Q methods plus the no-model
baseline, all on the augmented state (cell, phase):

  * Q-learning (K=0)            : no imagination -- the floor.
  * Tabular Dyna-Q (K=10)       : last-(s',r) lookup keyed on the augmented state.
  * One-hot WM Dyna-Q (K=10)    : neural WM whose input is the cell index only --
                                  structurally PHASE-BLIND, cannot see the walls.
  * Improved WM Dyna-Q (K=10)   : neural WM whose input adds the current 3x3 view --
                                  PHASE-AWARE, can see where the bar is now.

Both world models are RE-TRAINED on dynamic-obstacle data here (not the static
ones): on dynamic data the cell one-hot alone is no longer sufficient (the next
cell depends on the phase), which is exactly what forces the improved model to
read its window. Two figures are produced:

  1. figure_dynamic_prediction.png  -- next-cell prediction accuracy (the mechanism)
  2. figure_dynamic_learning_curves.png -- greedy success vs real steps (the headline)

NOTE (academic-integrity): this script is tooling. Whether the static-map ranking
reverses, by how much, and why -- the written interpretation -- is the author's.
"""

import os
import time
from collections import Counter

import numpy as np
import torch

from gridworld import DynamicGridWorld
from utils import (N_STATES, N_ACTIONS, state_to_index, sa_to_vector,
                   obs_a_to_vector, state_to_obs, build_obs_table,
                   IMPROVED_INPUT_DIM, IMPROVED_OBS_DIM)
from world_model import (TransitionNet, RewardNet, train_transition_model,
                         train_reward_model, train_improved_world_model,
                         improved_adapters, RESULTS_DIR)
from dyna_q import value_iteration
from dyna_q_dynamic import (dyna_q_dynamic, q_learning_dynamic,
                            get_all_transitions_dynamic, evaluate_greedy_dynamic)

torch.set_num_threads(1)

SEEDS = [0, 1, 2]
N_STEPS = 50_000
RECORD_EVERY = 1000
EVAL_EPISODES = 20
N_DATA = 10_000


# --------------------------------------------------------------------------- #
# 1. Data collection on the dynamic env (one pass -> both encodings + metadata)
# --------------------------------------------------------------------------- #
def collect_dynamic(env, n_transitions=N_DATA, seed=0):
    """Random-policy rollout on the dynamic env.

    Returns a dict of aligned arrays: the 68-dim one-hot input (X68, phase-blind),
    the 77-dim improved input with the CURRENT-phase window (X77), the next
    observation for the improved target (Y_obs), reward, and the raw (cell, phase,
    action, next_cell) metadata used for the tabular/phase-blind majority baselines.
    """
    env.rng.seed(seed)
    X68 = np.zeros((n_transitions, N_STATES + N_ACTIONS), dtype=np.float32)
    X77 = np.zeros((n_transitions, IMPROVED_INPUT_DIM), dtype=np.float32)
    Y_obs = np.zeros((n_transitions, IMPROVED_OBS_DIM), dtype=np.float32)
    Y_rew = np.zeros((n_transitions, 1), dtype=np.float32)
    cell = np.zeros(n_transitions, dtype=np.int64)
    phase = np.zeros(n_transitions, dtype=np.int64)
    act = np.zeros(n_transitions, dtype=np.int64)
    ncell = np.zeros(n_transitions, dtype=np.int64)

    s = env.reset()
    for i in range(n_transitions):
        walls_now = env._walls_for_phase(env.phase)
        ph = env.phase
        a = env.rng.randint(0, N_ACTIONS - 1)
        s_next, r, done, info = env.step(a)
        walls_next = env._walls_for_phase(info["phase_next"])
        X68[i] = sa_to_vector(s, a)
        X77[i] = obs_a_to_vector(s, a, walls=walls_now, goal=env.GOAL)
        Y_obs[i] = state_to_obs(s_next, walls=walls_next, goal=env.GOAL)
        Y_rew[i] = r
        cell[i], phase[i], act[i] = state_to_index(s), ph, a
        ncell[i] = state_to_index(s_next)
        s = env.reset() if done else s_next
    return {"X68": X68, "X77": X77, "Y_obs": Y_obs, "Y_rew": Y_rew,
            "cell": cell, "phase": phase, "act": act, "ncell": ncell}


def split_idx(n, val_frac=0.2, seed=0):
    rng = np.random.default_rng(seed)
    perm = rng.permutation(n)
    n_val = int(round(n * val_frac))
    return perm[n_val:], perm[:n_val]   # train_idx, val_idx


# --------------------------------------------------------------------------- #
# 2. Train both world models on the dynamic data
# --------------------------------------------------------------------------- #
def train_dynamic_world_models(data, seed=0, epochs=20, save=True):
    """Train the phase-blind one-hot WM and the phase-aware improved WM."""
    tr, va = split_idx(len(data["X68"]), seed=seed)

    onehot_tr = {"X": torch.from_numpy(data["X68"][tr]),
                 "y_next": torch.from_numpy(data["ncell"][tr]),
                 "y_rew": torch.from_numpy(data["Y_rew"][tr])}
    onehot_va = {"X": torch.from_numpy(data["X68"][va]),
                 "y_next": torch.from_numpy(data["ncell"][va]),
                 "y_rew": torch.from_numpy(data["Y_rew"][va])}
    tnet, acc_hist = train_transition_model(onehot_tr, onehot_va, epochs=epochs, seed=seed)
    rnet, _ = train_reward_model(onehot_tr, onehot_va, epochs=epochs, seed=seed)

    imp_tr = {"X": torch.from_numpy(data["X77"][tr]),
              "Y_obs": torch.from_numpy(data["Y_obs"][tr]),
              "Y_rew": torch.from_numpy(data["Y_rew"][tr]),
              "ns_idx": torch.from_numpy(data["ncell"][tr])}
    imp_va = {"X": torch.from_numpy(data["X77"][va]),
              "Y_obs": torch.from_numpy(data["Y_obs"][va]),
              "Y_rew": torch.from_numpy(data["Y_rew"][va]),
              "ns_idx": torch.from_numpy(data["ncell"][va])}
    # p_drop=0: on dynamic data the one-hot alone is already insufficient, so the
    # model is forced to read the window without artificial dropout.
    imp_model, imp_hist = train_improved_world_model(imp_tr, imp_va, p_drop=0.0,
                                                     epochs=epochs, seed=seed)

    if save:
        torch.save(tnet.state_dict(), os.path.join(RESULTS_DIR, "transition_net_dynamic.pt"))
        torch.save(rnet.state_dict(), os.path.join(RESULTS_DIR, "reward_net_dynamic.pt"))
        torch.save(imp_model.state_dict(), os.path.join(RESULTS_DIR, "improved_wm_dynamic.pt"))

    return {"tnet": tnet, "rnet": rnet, "improved": imp_model,
            "val_idx": va, "acc_hist": acc_hist, "imp_hist": imp_hist}


# --------------------------------------------------------------------------- #
# 3. Prediction-accuracy comparison (the mechanism figure)
# --------------------------------------------------------------------------- #
def _majority_accuracy(data, tr, va, keys):
    """Most-frequent-next-cell predictor keyed on `keys` (tuple of array names)."""
    table = {}
    counters = {}
    for i in tr:
        k = tuple(int(data[name][i]) for name in keys)
        counters.setdefault(k, Counter())[int(data["ncell"][i])] += 1
    for k, c in counters.items():
        table[k] = c.most_common(1)[0][0]
    correct = 0
    for i in va:
        k = tuple(int(data[name][i]) for name in keys)
        pred = table.get(k, -1)
        correct += (pred == int(data["ncell"][i]))
    return correct / len(va)


def prediction_accuracy(data, models, save=True):
    """Compute + plot next-cell prediction accuracy for the four approaches."""
    from plotting import plot_dynamic_prediction

    tr, va = split_idx(len(data["X68"]))
    ncell_va = data["ncell"][va]

    tnet, imp = models["tnet"], models["improved"]
    tnet.eval(); imp.eval()
    with torch.no_grad():
        oh_pred = tnet(torch.from_numpy(data["X68"][va])).argmax(1).numpy()
        imp_logits, _, _ = imp(torch.from_numpy(data["X77"][va]))
        imp_pred = imp_logits.argmax(1).numpy()
    acc_onehot = float((oh_pred == ncell_va).mean())
    acc_improved = float((imp_pred == ncell_va).mean())

    acc_tab_aug = _majority_accuracy(data, tr, va, ("cell", "phase", "act"))
    acc_blind = _majority_accuracy(data, tr, va, ("cell", "act"))

    labels = ["One-hot WM\n(phase-blind)", "Improved WM\n(phase-aware)",
              "Tabular memory\n(cell,phase,a)", "Phase-blind\nceiling (cell,a)"]
    accs = [acc_onehot, acc_improved, acc_tab_aug, acc_blind]
    print("Prediction accuracy on dynamic data (val):")
    for l, a in zip(labels, accs):
        print(f"  {l.splitlines()[0]:<18s}: {a:.2%}")

    fig = None
    if save:
        fig = plot_dynamic_prediction(
            labels, accs,
            ceilings={"noise ceiling": 0.85},
            title="Dynamic obstacles: next-cell prediction accuracy (P=4 phases)")
        np.savez(os.path.join(RESULTS_DIR, "dynamic_prediction.npz"),
                 labels=np.array(labels), accs=np.array(accs))
    return {"labels": labels, "accs": accs, "fig": fig}


# --------------------------------------------------------------------------- #
# 4. Control: four conditions x seeds on the augmented state
# --------------------------------------------------------------------------- #
def make_improved_encoder(env):
    """Phase-aware (cell, phase, action) -> 77-dim improved input batch."""
    P = env.n_phases
    tables = np.stack([build_obs_table(walls=env._walls_for_phase(p), goal=env.GOAL)
                       for p in range(P)])  # (P, 64, 73)

    def encode(cells, phases, acts):
        cells = np.asarray(cells); phases = np.asarray(phases); acts = np.asarray(acts)
        m = len(cells)
        X = np.zeros((m, IMPROVED_INPUT_DIM), dtype=np.float32)
        X[:, :IMPROVED_OBS_DIM] = tables[phases, cells]
        X[np.arange(m), IMPROVED_OBS_DIM + acts] = 1.0
        return X
    return encode


def run_condition_dynamic(label, seeds, tnet, rnet, K, imagine, n_steps=N_STEPS,
                          sa_encoder=None, snapshot_steps=None):
    evals, snaps = [], []
    steps = None
    for seed in seeds:
        env = DynamicGridWorld(seed=seed)
        res = dyna_q_dynamic(env, tnet, rnet, K=K, imagine=imagine, n_steps=n_steps,
                             record_every=RECORD_EVERY, eval_episodes=EVAL_EPISODES,
                             seed=seed, sa_encoder=sa_encoder,
                             snapshot_steps=snapshot_steps)
        steps = res["steps"]
        evals.append(res["eval_success_rate"])
        if snapshot_steps is not None:
            snaps.append(res["snapshots"])
        print(f"  [{label}] seed={seed}: final greedy success={res['eval_success_rate'][-1]:.2%}")
    out = {"label": label, "steps": steps, "eval_success": np.array(evals)}
    if snapshot_steps is not None:
        out["snapshots"] = snaps
    return out


def main(seeds=SEEDS, n_steps=N_STEPS, save=True):
    os.makedirs(RESULTS_DIR, exist_ok=True)
    t0 = time.time()

    # --- visualization-first: phase layouts + a random rollout ---------------
    from plotting import plot_dynamic_phases, animate_dynamic_episode
    env0 = DynamicGridWorld(seed=0)
    plot_dynamic_phases(env0)
    animate_dynamic_episode(DynamicGridWorld(seed=3), seed=3)

    # --- VI ceiling on the augmented MDP ------------------------------------
    P_aug, R_aug = get_all_transitions_dynamic(env0)
    _, Q_vi, _ = value_iteration(P_aug, R_aug)
    vi_success, _ = evaluate_greedy_dynamic(Q_vi, env0, n_episodes=200)
    print(f"[VI augmented] greedy success={vi_success:.2%}")

    # --- train both world models on dynamic data ----------------------------
    print("Collecting dynamic data + training world models ...")
    data = collect_dynamic(DynamicGridWorld(seed=0), seed=0)
    models = train_dynamic_world_models(data, seed=0)
    pred = prediction_accuracy(data, models, save=save)

    imp_t, imp_r = improved_adapters(models["improved"])
    enc = make_improved_encoder(env0)

    # --- four control conditions x seeds ------------------------------------
    conds = [
        ("Q-learning (K=0)",            None,         None,         0,  "tabular", None, "#1f77b4", "k0"),
        ("Tabular Dyna-Q (K=10)",       None,         None,         10, "tabular", None, "#ff7f0e", "tabular"),
        ("One-hot WM Dyna-Q (K=10)",    models["tnet"], models["rnet"], 10, "neural", None, "#d62728", "onehot"),
        ("Improved WM Dyna-Q (K=10)",   imp_t,        imp_r,        10, "neural", enc,  "#2ca02c", "improved"),
    ]
    results = []
    for label, tn, rn, K, imagine, encoder, _, _ in conds:
        print(f"Running {label} x{len(seeds)} seeds ...")
        results.append(run_condition_dynamic(label, seeds, tn, rn, K, imagine,
                                             n_steps=n_steps, sa_encoder=encoder))

    if save:
        payload = {"steps": results[0]["steps"], "vi_success": vi_success,
                   "seeds": np.array(seeds), "pred_accs": np.array(pred["accs"])}
        for (label, *_, key), res in zip(conds, results):
            payload[f"{key}_eval"] = res["eval_success"]
        np.savez(os.path.join(RESULTS_DIR, "dynamic_curves.npz"), **payload)
        print(f"Saved dynamic curves to {RESULTS_DIR}/dynamic_curves.npz")

    print(f"Total dynamic experiment time: {time.time() - t0:.1f}s")
    return {"conds": conds, "results": results, "vi_success": vi_success,
            "pred": pred, "models": models}


if __name__ == "__main__":
    from plotting import plot_learning_curves

    out = main()
    curves = [{"label": res["label"], "steps": res["steps"],
               "data": res["eval_success"], "color": color}
              for (label, *_, color, _), res in zip(out["conds"], out["results"])]
    fig = plot_learning_curves(
        curves, vi_ceiling=out["vi_success"],
        title="Dynamic obstacles: greedy success rate vs real steps (augmented state)",
        save_path=os.path.join(RESULTS_DIR, "figure_dynamic_learning_curves.png"),
    )
    print(f"Saved {fig}")