dsh-eval-guard 0.7.4 → 0.7.5

package/assets/neural-pendulum/rl-feasibility/evaluate.py

@@ -1,107 +0,0 @@
-"""Deterministic evaluation at the fixed initial state (and optional starts).
-
-The evaluation uses the simulator API exactly like training (angles only), the
-same observation builder / obs mode, and the same trainer-side reward.  There
-is no action noise.
-
-    default fixed state: theta = [0.013, -0.008, 0.021] rad, omega = 0
-    optional starts    : drawn from the training reset distribution
-
-The episode always runs the full 3000 steps so that t_can, ever_up, return and
-the final angle error are comparable across algorithms.
-"""
-from __future__ import annotations
-
-import numpy as np
-import torch
-
-from obs import DT, make_numpy_obs, wrap_to_pi
-from sim import N_ACT, N_LINKS, UP_TOL, PendulumSim
-from task import EVAL_OMEGA0, EVAL_THETA0, NumpyTask
-
-
-def evaluate(policy_fn, p, episode_steps=3000, reward_w=None, obs_mode="full"):
-    """Run one full deterministic episode at the fixed initial state.
-
-    policy_fn(obs: np.ndarray[obs_dim]) -> np.ndarray[2]
-    """
-    sim = PendulumSim(p, episode_steps=episode_steps)
-    task = NumpyTask(p, reward_w=reward_w, obs_mode=obs_mode)
-    theta = sim.reset(EVAL_THETA0, EVAL_OMEGA0)
-    obs = task.reset(theta)
-    acts = np.zeros((episode_steps, N_ACT), dtype=np.float64)
-    for k in range(episode_steps):
-        u = np.clip(np.asarray(policy_fn(obs), dtype=np.float64).reshape(N_ACT),
-                    -1.0, 1.0)
-        acts[k] = u
-        theta = sim.step(u)
-        obs, _ = task.step(theta, u)
-    summ = task.tracker.summary()
-    a = wrap_to_pi(np.asarray(theta, dtype=np.float64) - np.pi)
-    return dict(t_can=float(summ["t_can"]),
-                success=float(summ["t_can"] >= 10.0),
-                ever_up=float(summ["ever_up"]),
-                first_up_t=float(summ["first_up_t"]),
-                ep_return=float(summ["ep_return"]),
-                final_ang_err=float(np.abs(a).max()),
-                up_frac=float(summ["up_frac"]),
-                mean_abs_action=float(np.mean(np.abs(acts))),
-                action_std=float(np.std(acts)),
-                mean_u1=float(acts[:, 0].mean()),
-                mean_u3=float(acts[:, 1].mean()),
-                steps=int(episode_steps))
-
-
-def evaluate_starts(policy_fn, p, starts, episode_steps=3000, reward_w=None,
-                    obs_mode="full"):
-    """Deterministic rollouts from an explicit list of (theta, omega) starts."""
-    rows = []
-    for theta0, omega0 in starts:
-        sim = PendulumSim(p, episode_steps=episode_steps)
-        task = NumpyTask(p, reward_w=reward_w, obs_mode=obs_mode)
-        theta = sim.reset(theta0, omega0)
-        obs = task.reset(theta)
-        for _ in range(episode_steps):
-            u = np.clip(np.asarray(policy_fn(obs), dtype=np.float64).reshape(N_ACT),
-                        -1.0, 1.0)
-            theta = sim.step(u)
-            obs, _ = task.step(theta, u)
-        summ = task.tracker.summary()
-        a = wrap_to_pi(np.asarray(theta, dtype=np.float64) - np.pi)
-        rows.append(dict(t_can=float(summ["t_can"]),
-                         success=float(summ["t_can"] >= 10.0),
-                         ever_up=float(summ["ever_up"]),
-                         ep_return=float(summ["ep_return"]),
-                         final_ang_err=float(np.abs(a).max())))
-    tc = np.array([r["t_can"] for r in rows], dtype=np.float64)
-    su = np.array([r["success"] for r in rows], dtype=np.float64)
-    re = np.array([r["ep_return"] for r in rows], dtype=np.float64)
-    ae = np.array([r["final_ang_err"] for r in rows], dtype=np.float64)
-    eu = np.array([r["ever_up"] for r in rows], dtype=np.float64)
-    return dict(n=int(len(rows)),
-                success_rate=float(su.mean()) if len(rows) else 0.0,
-                ever_up_rate=float(eu.mean()) if len(rows) else 0.0,
-                mean_t_can=float(tc.mean()) if len(rows) else 0.0,
-                median_t_can=float(np.median(tc)) if len(rows) else 0.0,
-                min_t_can=float(tc.min()) if len(rows) else 0.0,
-                max_t_can=float(tc.max()) if len(rows) else 0.0,
-                mean_return=float(re.mean()) if len(rows) else 0.0,
-                mean_final_ang_err=float(ae.mean()) if len(rows) else 0.0,
-                rows=rows)
-
-
-def make_ppo_policy(actor):
-    def policy(obs):
-        with torch.no_grad():
-            x = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)
-            return actor(x).squeeze(0).numpy()
-    return policy
-
-
-def make_sac_policy(actor):
-    def policy(obs):
-        with torch.no_grad():
-            x = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)
-            a, _ = actor.sample(x, det=True)
-            return a.squeeze(0).numpy()
-    return policy

package/assets/neural-pendulum/rl-feasibility/finalize_run.py

@@ -1,168 +0,0 @@
-"""Reconstruct results/<exp>.json for a finished run from its checkpoints,
-TensorBoard event file and log.  Useful when the trainer crashed after saving
-checkpoints but before writing the JSON (e.g. the PPO robustness bug).
-
-Usage:
-  python3 finalize_run.py --exp ppo_full_base_rand --algo ppo \
-      --obs full --reward base --init rand --envs 64 --minutes 45
-
-The fixed-state eval and the 8-start robustness check are re-run here with the
-same code as the trainers, so the numbers are real, not copied from the log.
-The eval curve is parsed from logs/<exp>.log and the wall time from TB.
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import re
-from pathlib import Path
-
-import torch
-
-from evaluate import (evaluate, evaluate_starts, make_ppo_policy,
-                      make_sac_policy)
-from models import ActorCritic, SquashedGaussianActor
-from obs import obs_spec
-from sim import EPISODE_STEPS, Params
-from task import get_reward_preset, sample_train_init
-
-
-def build_policy(ck, algo, obs_dim):
-    a = ck.get("args") or {}
-    hidden = int(a.get("hidden", 128 if algo == "ppo" else 256))
-    depth = int(a.get("depth", 2))
-    if algo == "ppo":
-        net = ActorCritic(obs_dim, act_dim=2, hidden=hidden, depth=depth)
-        net.load_state_dict(ck["model"])
-        return make_ppo_policy(net.actor), net
-    net = SquashedGaussianActor(obs_dim, 2, hidden=hidden, depth=depth)
-    net.load_state_dict(ck["model"])
-    return make_sac_policy(net), net
-
-
-def parse_curve(log_path):
-    rows = []
-    if not os.path.exists(log_path):
-        return rows
-    pat = re.compile(
-        r"iter\s+(\d+)\s+steps\s+([\d.]+)M.*?EVAL tc\s+([\d.]+)\s+"
-        r"succ\s+(\d+)\s+everup\s+(\d+)\s+angerr\s+([\d.]+)\s+"
-        r"ret\s+([-\d.]+)")
-    pat_sac = re.compile(
-        r"steps\s+([\d.]+)k\s+upd\s+(\d+).*?EVAL tc\s+([\d.]+)\s+"
-        r"succ\s+(\d+)\s+everup\s+(\d+)\s+angerr\s+([\d.]+)\s+"
-        r"ret\s+([-\d.]+)")
-    for line in open(log_path, errors="ignore"):
-        m = pat.search(line)
-        if m:
-            rows.append(dict(iter=int(m.group(1)),
-                             global_steps=int(float(m.group(2)) * 1e6),
-                             t_can=float(m.group(3)),
-                             success=float(m.group(4)),
-                             ever_up=float(m.group(5)),
-                             final_ang_err=float(m.group(6)),
-                             ep_return=float(m.group(7))))
-            continue
-        m = pat_sac.search(line)
-        if m:
-            rows.append(dict(global_steps=int(float(m.group(1)) * 1e3),
-                             n_updates=int(m.group(2)),
-                             t_can=float(m.group(3)),
-                             success=float(m.group(4)),
-                             ever_up=float(m.group(5)),
-                             final_ang_err=float(m.group(6)),
-                             ep_return=float(m.group(7))))
-    return rows
-
-
-def tb_last(tb_root, exp, tag):
-    from tensorboard.backend.event_processing.event_accumulator import (
-        EventAccumulator)
-    d = Path(tb_root) / exp
-    if not d.is_dir():
-        return None
-    ea = EventAccumulator(str(d), size_guidance={"scalars": 0})
-    ea.Reload()
-    if tag not in ea.Tags().get("scalars", []):
-        return None
-    return ea.Scalars(tag)[-1].value
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--exp", required=True)
-    ap.add_argument("--algo", required=True, choices=["ppo", "sac"])
-    ap.add_argument("--obs", required=True)
-    ap.add_argument("--reward", required=True)
-    ap.add_argument("--init", required=True)
-    ap.add_argument("--envs", type=int, required=True)
-    ap.add_argument("--minutes", type=float, required=True)
-    ap.add_argument("--seed", type=int, default=0)
-    ap.add_argument("--ckpt-dir", default="checkpoints")
-    ap.add_argument("--tb-root", default="tb")
-    ap.add_argument("--out", default="results")
-    ap.add_argument("--log-dir", default="logs")
-    ap.add_argument("--robust-starts", type=int, default=8)
-    args = ap.parse_args()
-
-    ck_final = torch.load(os.path.join(args.ckpt_dir, args.exp + "_final.pt"),
-                          map_location="cpu")
-    ck_best = torch.load(os.path.join(args.ckpt_dir, args.exp + "_best.pt"),
-                         map_location="cpu")
-    obs_dim = ck_final["obs_dim"]
-    spec = ck_final.get("obs") or obs_spec(args.obs)
-    reward_w = ck_final.get("reward") or get_reward_preset(args.reward)
-    p = Params()
-    pol_final, _ = build_policy(ck_final, args.algo, obs_dim)
-    pol_best, _ = build_policy(ck_best, args.algo, obs_dim)
-    ev_final = evaluate(pol_final, p, EPISODE_STEPS, reward_w=reward_w,
-                        obs_mode=args.obs)
-    ev_best = evaluate(pol_best, p, EPISODE_STEPS, reward_w=reward_w,
-                       obs_mode=args.obs)
-
-    g = torch.Generator().manual_seed(args.seed * 1000 + 123)
-    th, om = sample_train_init(args.robust_starts, gen=g)
-    starts = list(zip(th.numpy(), om.numpy()))
-    rob = evaluate_starts(pol_final, p, starts, EPISODE_STEPS,
-                          reward_w=reward_w, obs_mode=args.obs)
-    rob_best = evaluate_starts(pol_best, p, starts, EPISODE_STEPS,
-                               reward_w=reward_w, obs_mode=args.obs)
-
-    curve = parse_curve(os.path.join(args.log_dir, args.exp + ".log"))
-    elapsed = tb_last(args.tb_root, args.exp, "time/elapsed_min")
-    wall = float(elapsed) * 60.0 if elapsed is not None else 0.0
-    steps = ck_final.get("global_steps")
-    if not steps:
-        steps = tb_last(args.tb_root, args.exp, "train/global_steps")
-    out = dict(algo=args.algo, exp=args.exp, seed=args.seed,
-               envs=args.envs, obs=spec, obs_dim=obs_dim,
-               init=ck_final.get("init"), reward=reward_w,
-               reward_preset=args.reward, args=ck_final.get("args"),
-               minutes_budget=args.minutes, global_steps=int(steps or 0),
-               wall_time_s=wall,
-               steps_per_s=(float(steps) / wall) if wall > 0 else None,
-               final_eval=ev_final, best_eval=ev_best,
-               best_t_can=ev_best["t_can"], eval_curve=curve,
-               robust=rob, robust_best=rob_best,
-               checkpoints=dict(best=args.exp + "_best.pt",
-                                last=args.exp + "_last.pt",
-                                final=args.exp + "_final.pt"),
-               reconstructed=True,
-               reconstruct_note=("results JSON reconstructed by finalize_run.py "
-                                 "after the trainer crashed post-checkpoint; "
-                                 "fixed eval and robustness were re-run."))
-    with open(os.path.join(args.out, args.exp + ".json"), "w") as f:
-        json.dump(out, f, indent=2)
-    print(json.dumps(dict(exp=args.exp, steps=out["global_steps"],
-                          wall_min=round(wall / 60.0, 1),
-                          final_t_can=ev_final["t_can"],
-                          best_t_can=ev_best["t_can"],
-                          final_success=ev_final["success"],
-                          best_success=ev_best["success"],
-                          robust_final=rob["success_rate"],
-                          robust_best=rob_best["success_rate"]), indent=2))
-
-
-if __name__ == "__main__":
-    main()

package/assets/neural-pendulum/rl-feasibility/models.py

@@ -1,124 +0,0 @@
-"""Networks: PPO actor-critic MLP, SAC squashed-Gaussian actor, twin Q.
-
-PPO: 2x128 tanh, separate actor and critic, state-independent log_std capped
-to [-2.3, -0.25] (the recipe of the sibling rounds).
-SAC: 2x256 tanh actor with tanh-Gaussian head and a shared-trunk twin Q (two
-output heads).  Sharing the critic trunk keeps the CPU update affordable while
-preserving the double-Q target.
-"""
-from __future__ import annotations
-
-import math
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-
-class MLP(nn.Module):
-    def __init__(self, in_dim, hidden=128, depth=2, out_dim=1, out_gain=0.01):
-        super().__init__()
-        layers = []
-        d = in_dim
-        for _ in range(depth):
-            lin = nn.Linear(d, hidden)
-            nn.init.orthogonal_(lin.weight, gain=np.sqrt(2))
-            nn.init.zeros_(lin.bias)
-            layers += [lin, nn.Tanh()]
-            d = hidden
-        self.body = nn.Sequential(*layers)
-        self.out = nn.Linear(d, out_dim)
-        nn.init.orthogonal_(self.out.weight, gain=out_gain)
-        nn.init.zeros_(self.out.bias)
-
-    def forward(self, x):
-        return self.out(self.body(x))
-
-
-# ------------------------------- PPO ---------------------------------------
-LOG_STD_MIN, LOG_STD_MAX = -2.3, -0.25
-
-
-class ActorCritic(nn.Module):
-    def __init__(self, obs_dim, act_dim=2, hidden=128, depth=2):
-        super().__init__()
-        self.act_dim = act_dim
-        self.actor = MLP(obs_dim, hidden, depth, act_dim)
-        self.critic = MLP(obs_dim, hidden, depth, 1, out_gain=1.0)
-        self.log_std = nn.Parameter(torch.full((act_dim,), -0.7))
-
-    def dist(self, obs):
-        mean = self.actor(obs)
-        std = self.effective_log_std().exp().expand_as(mean)
-        return torch.distributions.Normal(mean, std)
-
-    def effective_log_std(self):
-        return self.log_std.clamp(LOG_STD_MIN, LOG_STD_MAX)
-
-    def value(self, obs):
-        return self.critic(obs).squeeze(-1)
-
-    def deterministic_action(self, obs):
-        return self.actor(obs)
-
-
-# ------------------------------- SAC ---------------------------------------
-SAC_LOG_STD_MIN, SAC_LOG_STD_MAX = -5.0, 2.0
-
-
-class SquashedGaussianActor(nn.Module):
-    def __init__(self, in_dim, act_dim=2, hidden=256, depth=2):
-        super().__init__()
-        self.body = nn.Sequential()
-        d = in_dim
-        for _ in range(depth):
-            lin = nn.Linear(d, hidden)
-            nn.init.orthogonal_(lin.weight, gain=math.sqrt(2))
-            nn.init.zeros_(lin.bias)
-            self.body.extend([lin, nn.Tanh()])
-            d = hidden
-        self.mean = nn.Linear(hidden, act_dim)
-        self.log_std = nn.Linear(hidden, act_dim)
-        nn.init.orthogonal_(self.mean.weight, gain=0.01)
-        nn.init.zeros_(self.mean.bias)
-        nn.init.orthogonal_(self.log_std.weight, gain=0.01)
-        nn.init.zeros_(self.log_std.bias)
-        self.act_dim = act_dim
-
-    def forward(self, obs):
-        h = self.body(obs)
-        return self.mean(h), self.log_std(h).clamp(SAC_LOG_STD_MIN, SAC_LOG_STD_MAX)
-
-    def sample(self, obs, det=False):
-        if det:
-            return torch.tanh(self.mean(self.body(obs))), None
-        mean, log_std = self(obs)
-        std = log_std.exp()
-        eps = torch.randn_like(mean)
-        x = mean + std * eps
-        a = torch.tanh(x)
-        logp = (-0.5 * eps ** 2 - log_std - 0.5 * math.log(2 * math.pi)).sum(-1)
-        logp = logp - torch.log(1 - a ** 2 + 1e-6).sum(-1)
-        return a, logp
-
-
-class TwinQ(nn.Module):
-    """Shared trunk, two Q heads; forward returns (B, 2)."""
-
-    def __init__(self, in_dim, act_dim=2, hidden=256, depth=2):
-        super().__init__()
-        self.body = nn.Sequential()
-        d = in_dim + act_dim
-        for _ in range(depth):
-            lin = nn.Linear(d, hidden)
-            nn.init.orthogonal_(lin.weight, gain=math.sqrt(2))
-            nn.init.zeros_(lin.bias)
-            self.body.extend([lin, nn.Tanh()])
-            d = hidden
-        self.q = nn.Linear(hidden, 2)
-        nn.init.orthogonal_(self.q.weight, gain=1.0)
-        nn.init.zeros_(self.q.bias)
-
-    def forward(self, obs, act):
-        x = torch.cat([obs, act], dim=-1)
-        return self.q(self.body(x))

package/assets/neural-pendulum/rl-feasibility/obs.py

@@ -1,205 +0,0 @@
-"""Policy-side angle observation for the 3-link pendulum (switchable modes).
-
-The simulator only exposes the three joint angles theta.  The observation is
-built entirely from that angle stream (no omega, no tip, no energy):
-
-    a_i(t)  = wrap_to_pi(theta_i(t) - pi)                  (upright-relative)
-    da_i(t) = wrap_to_pi(a_i(t) - a_i(t-1)) / DT,          da_i(0) = 0 at reset
-
-    frame(t) = [cos a1, sin a1, cos a2, sin a2, cos a3, sin a3,
-                da1/5, da2/5, da3/5]                        (9 dims, diff=True)
-    frame(t) = [cos a1, sin a1, cos a2, sin a2, cos a3, sin a3]
-                                                            (6 dims, diff=False)
-
-    obs(t)   = H frames, oldest first, flattened
-
-Modes (--obs):
-    full      H=4, diff=True   -> 36 dims  (the original observation)
-    angle     H=4, diff=False  -> 24 dims  (cos/sin only, no manual velocity)
-    full_h1   H=1, diff=True   ->  9 dims  (single frame incl. da/5)
-    angle_h1  H=1, diff=False  ->  6 dims  (single frame cos/sin only)
-    full_h2   H=2, diff=True   -> 18 dims  (extra mode, optional)
-    angle_h2  H=2, diff=False  -> 12 dims  (extra mode, optional)
-
-Even when diff=False the builder still computes da internally, because the
-trainer-side reward uses it; the channel is only left out of the policy
-observation.  theta is used in radians exactly as produced by the simulator
-(never divided by pi).  The history is initialised by repeating the reset
-frame H times.
-"""
-from __future__ import annotations
-
-import numpy as np
-
-DT = 0.01
-H = 4
-VEL_SCALE = 5.0
-
-# mode -> (history length, whether the manual-difference velocity is observed)
-OBS_MODES = {
-    "full": dict(h=4, diff=True),
-    "angle": dict(h=4, diff=False),
-    "full_h1": dict(h=1, diff=True),
-    "angle_h1": dict(h=1, diff=False),
-    "full_h2": dict(h=2, diff=True),
-    "angle_h2": dict(h=2, diff=False),
-    "angle_h8": dict(h=8, diff=False),   # extra: can longer history replace da?
-}
-
-
-def frame_dim(diff: bool) -> int:
-    return 9 if diff else 6
-
-
-def obs_spec(mode: str = "full") -> dict:
-    if mode not in OBS_MODES:
-        raise ValueError(f"unknown obs mode {mode!r}; choose from {sorted(OBS_MODES)}")
-    c = OBS_MODES[mode]
-    fd = frame_dim(c["diff"])
-    notes = ("cos/sin(a) + da/5, H=%d, a=wrap(theta-pi), da is the manual "
-             "finite difference, da(0)=0" % c["h"])
-    if not c["diff"]:
-        notes = ("cos/sin(a) only (NO velocity channel), H=%d, "
-                 "a=wrap(theta-pi)" % c["h"])
-    return dict(mode=mode, enc="sincos", diff=bool(c["diff"]), h=int(c["h"]),
-                vel_scale=VEL_SCALE, dt=DT, frame_dim=fd,
-                dim=int(c["h"] * fd), notes=notes)
-
-
-# Backwards-compatible constants for the original full observation.
-FRAME_DIM = frame_dim(True)
-OBS_DIM = obs_spec("full")["dim"]
-
-
-def wrap_to_pi(a):
-    return (a + np.pi) % (2.0 * np.pi) - np.pi
-
-
-def frame_np(a, da, diff=True):
-    a = np.asarray(a, dtype=np.float64)
-    cs = np.empty(6, dtype=np.float64)
-    cs[0::2] = np.cos(a)
-    cs[1::2] = np.sin(a)
-    if not diff:
-        return cs
-    return np.concatenate([cs, np.asarray(da, dtype=np.float64) / VEL_SCALE])
-
-
-class NumpyAngleObs:
-    """Float64 single-env builder.  reset/update take theta only."""
-
-    def __init__(self, h: int = H, vel_scale: float = VEL_SCALE, dt: float = DT,
-                 diff: bool = True):
-        self.h = int(h)
-        self.diff = bool(diff)
-        self.fdim = frame_dim(self.diff)
-        self.vel_scale = float(vel_scale)
-        self.dt = float(dt)
-        self.a_prev = np.zeros(3)
-        self.da = np.zeros(3)
-        self.frames = np.zeros((self.h, self.fdim))
-        self._obs = np.zeros(self.h * self.fdim, dtype=np.float32)
-
-    def _frame(self, a, da):
-        cs = np.empty(6, dtype=np.float64)
-        cs[0::2] = np.cos(a)
-        cs[1::2] = np.sin(a)
-        if not self.diff:
-            return cs
-        return np.concatenate([cs, da / self.vel_scale])
-
-    def _flatten(self):
-        self._obs = self.frames.reshape(-1).astype(np.float32)
-        return self._obs
-
-    def reset(self, theta):
-        a = wrap_to_pi(np.asarray(theta, dtype=np.float64).reshape(3) - np.pi)
-        self.a_prev = a.copy()
-        self.da = np.zeros(3)
-        f0 = self._frame(a, self.da)
-        self.frames = np.tile(f0, (self.h, 1))
-        return self._flatten()
-
-    @property
-    def obs(self):
-        return self._obs
-
-    def update(self, theta):
-        a = wrap_to_pi(np.asarray(theta, dtype=np.float64).reshape(3) - np.pi)
-        self.da = wrap_to_pi(a - self.a_prev) / self.dt
-        self.a_prev = a.copy()
-        f = self._frame(a, self.da)
-        if self.h > 1:
-            self.frames = np.roll(self.frames, -1, axis=0)
-        self.frames[-1] = f
-        return self._flatten()
-
-
-class TorchAngleObs:
-    """Batched float32 builder.  reset/update take theta only."""
-
-    def __init__(self, n: int, h: int = H, vel_scale: float = VEL_SCALE,
-                 dt: float = DT, device="cpu", diff: bool = True):
-        import torch
-        self.torch = torch
-        self.n = int(n)
-        self.h = int(h)
-        self.diff = bool(diff)
-        self.fdim = frame_dim(self.diff)
-        self.vel_scale = float(vel_scale)
-        self.dt = float(dt)
-        self.device = device
-        self.a_prev = torch.zeros(n, 3, device=device)
-        self.da = torch.zeros(n, 3, device=device)
-        self.frames = torch.zeros(n, self.h, self.fdim, device=device)
-
-    def _frame(self, a, da):
-        t = self.torch
-        cs = t.stack([t.cos(a), t.sin(a)], dim=2).reshape(a.shape[0], 6)
-        if not self.diff:
-            return cs
-        return t.cat([cs, da / self.vel_scale], dim=1)
-
-    def obs(self):
-        return self.frames.reshape(self.n, -1)
-
-    def reset(self, theta, idx=None):
-        t = self.torch
-        theta = t.as_tensor(theta, dtype=t.float32, device=self.device).reshape(-1, 3)
-        if idx is None:
-            a = t.remainder(theta - np.pi + np.pi, 2 * np.pi) - np.pi
-            self.a_prev = a.clone()
-            self.da = t.zeros_like(a)
-            f0 = self._frame(a, self.da)
-            self.frames = f0.unsqueeze(1).expand(-1, self.h, -1).clone()
-            return self.obs()
-        idx = t.as_tensor(idx, dtype=t.long, device=self.device).flatten()
-        a = t.remainder(theta - np.pi + np.pi, 2 * np.pi) - np.pi
-        self.a_prev[idx] = a
-        self.da[idx] = 0.0
-        f0 = self._frame(a, t.zeros_like(a))
-        self.frames[idx] = f0.unsqueeze(1).expand(-1, self.h, -1)
-        return self.obs()
-
-    def update(self, theta):
-        t = self.torch
-        theta = t.as_tensor(theta, dtype=t.float32, device=self.device).reshape(-1, 3)
-        a = t.remainder(theta - np.pi + np.pi, 2 * np.pi) - np.pi
-        self.da = (t.remainder(a - self.a_prev + np.pi, 2 * np.pi) - np.pi) / self.dt
-        self.a_prev = a.clone()
-        f = self._frame(a, self.da)
-        if self.h > 1:
-            self.frames = t.cat([self.frames[:, 1:], f.unsqueeze(1)], dim=1)
-        else:
-            self.frames[:, -1] = f
-        return self.obs()
-
-
-def make_numpy_obs(mode: str = "full", **kw):
-    c = OBS_MODES[mode]
-    return NumpyAngleObs(h=c["h"], diff=c["diff"], **kw)
-
-
-def make_torch_obs(n: int, mode: str = "full", **kw):
-    c = OBS_MODES[mode]
-    return TorchAngleObs(n, h=c["h"], diff=c["diff"], **kw)

package/assets/neural-pendulum/rl-feasibility/plot_selected.py

@@ -1,55 +0,0 @@
-"""Plot selected eval t_can curves (PPO left, SAC right)."""
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-ROOT = Path(__file__).resolve().parent
-RES = ROOT / "results"
-
-PPO = ["ppo_full_base_rand", "ppo_full_base_rand_s1",
-       "ppo_full_base_rand_s2", "ppo_angle_base_rand", "ppo_angle_h8_rand",
-       "ppo_full_h1_rand", "ppo_full_h1_mix", "ppo_full_base_fixed"]
-SAC = ["sac_full_base_rand", "sac_full_base_fixed", "sac_angle_base_rand"]
-
-
-def load(exp):
-    p = RES / (exp + ".json")
-    if not p.exists():
-        return None
-    with open(p) as f:
-        return json.load(f)
-
-
-def main():
-    fig, axes = plt.subplots(1, 2, figsize=(14, 5.0), sharey=True)
-    for ax, names, title in [(axes[0], PPO, "PPO"), (axes[1], SAC, "SAC")]:
-        for exp in names:
-            r = load(exp)
-            if not r:
-                continue
-            c = r.get("eval_curve") or []
-            if not c:
-                continue
-            x = [v["global_steps"] / 1e6 for v in c]
-            y = [v["t_can"] for v in c]
-            ax.plot(x, y, lw=1.3, label=f"{exp} (best {r['best_t_can']:.2f}s)")
-        ax.axhline(10.0, color="k", ls=":", lw=1)
-        ax.set_title(f"{title}: fixed-initial eval t_can")
-        ax.set_xlabel("environment steps (M)")
-        ax.set_ylabel("t_can (s)")
-        ax.grid(alpha=0.3)
-        ax.legend(fontsize=7)
-    fig.tight_layout()
-    out = RES / "curves_selected.png"
-    fig.savefig(out, dpi=130)
-    plt.close(fig)
-    print(out)
-
-
-if __name__ == "__main__":
-    main()