File size: 3,020 Bytes
872cf4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Recover the per-q validation distance profile for checkpoints that predate
the profile logging in ``train_controller.py``.
The original controller was trained before ``evaluate()`` recorded
``val_profile``, so its arrival-timing signature — the thing Figure 2 shows —
is missing from its checkpoint while every later run has it. This recomputes
it with the same function, the same held-out split, and the same batch order,
and writes it next to the diagnostics so the figure script can read it.
python report/recover_profiles.py --controllers data/runs/controller/controller.pt
"""
import argparse
import json
import sys
from pathlib import Path
import torch
from torch.utils.data import DataLoader
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from lejepa_control.data import LatentGoalDataset, split_episodes # noqa: E402
from lejepa_control.solver import load_controller # noqa: E402
from lejepa_control.world_model import load_lewm # noqa: E402
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'scripts'))
from train_controller import evaluate # noqa: E402
def main():
p = argparse.ArgumentParser()
p.add_argument('--controllers', nargs='+',
default=['data/runs/controller/controller.pt'])
p.add_argument('--latents', default='data/latents')
p.add_argument('--batch-size', type=int, default=128)
p.add_argument('--out', default='data/runs/diagnostics/profiles_posthoc.json')
args = p.parse_args()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
stats = json.loads((Path(args.latents) / 'stats.json').read_text())
model = load_lewm(device=device)
_, val_eps = split_episodes(stats['n_episodes'])
val_set = LatentGoalDataset(
args.latents, max_offset=5, episodes=val_eps, horizon=5
)
out = {}
for path in args.controllers:
name = Path(path).parent.name
controller, ckpt = load_controller(
path, latent_dim=stats['latent_dim'], device=device
)
# same batches for every checkpoint, so profiles are comparable
loader = DataLoader(
val_set, batch_size=args.batch_size, shuffle=True,
generator=torch.Generator().manual_seed(0)
)
val = evaluate(controller, model, loader, device)
out[name] = {
'step': int(ckpt['step']),
'terminal': val['terminal'],
'first': val['first'],
'arrival': val['arrival'],
'profile': val['profile'],
}
print(f'{name}: step {ckpt["step"]} terminal {val["terminal"]:.4f} '
f'arrival {val["arrival"]:.4f}')
for q, prof in val['profile'].items():
print(f' q={q}: {prof}')
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
Path(args.out).write_text(json.dumps(out, indent=2))
print(f'\nwrote {args.out}')
if __name__ == '__main__':
main()
|