Buckets:
| """cwm3/train/train_cwm3_rssm.py — the joint trainer for the RSSM decoder variant (GPU-baselines experiment E), train/train_cwm3_joint.py used AS A LIBRARY | |
| (that file is owned by the final-assembly agent and is not edited). Same CLI as train_cwm3_joint.py (--stage F1 --variant rssm ...). | |
| Two process-local hooks on the library module (numerically no-ops for any model without RSSM decoders): | |
| * J.joint_loss -> the original loss + lambda_KL (cfg lambda `KL`, default 1.0) * sum of the RSSM decoders' free-bits KL terms of the current batch | |
| (part `kl_rssm`; NOT a chain_* part, so the selection objective / early stop are unchanged); | |
| * J.prep_joint -> the original batch prep, plus decoders_rssm.CONTEXT["tg"] = the batch targets (teacher-forcing posterior input) — read by the | |
| decoders ONLY in training mode; cleared when the batch has no targets. val_objective / predict_split run the model in eval mode, | |
| where the decoders draw 32 prior samples and never touch CONTEXT (tests/test_rssm_decoder.py). | |
| Everything else (base loading, checkpoints, prediction packs, layouts, run json) is train_cwm3_joint.py unchanged. No openpilot DMS field is read. | |
| python train/train_cwm3_rssm.py --stage F1 --variant rssm --seed 0 --no-lock --minutes 600 --pack-version v2 \\ | |
| --base-v can4s_stats2 --base-u kin_can_noK_clip8_ln --base-e pooled_P --run-suffix _final --base-dir-v ... --base-dir-u ... --base-dir-e ... --out results/runs/F1_rssm_s0_final | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_NUM_THREADS", "4") | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval", CWM3.parent / "s12_wm"): | |
| sys.path.insert(0, str(p)) | |
| import train_cwm3_joint as J # noqa: E402 | |
| from models import decoders_rssm as R # noqa: E402 | |
| _orig_loss, _orig_prep = J.joint_loss, J.prep_joint | |
| def joint_loss_rssm(out, tg, lam): | |
| total, parts = _orig_loss(out, tg, lam) | |
| kl = R.collect_kl() | |
| if kl is not None: | |
| parts["kl_rssm"] = kl.detach() | |
| total = total + float(lam.get("KL", 1.0)) * kl | |
| return total, parts | |
| def prep_joint_rssm(store, idx, st, pc, need_targets=True, region_keys=False): | |
| b, tg, vv = _orig_prep(store, idx, st, pc, need_targets, region_keys) | |
| R.CONTEXT["tg"] = tg if need_targets else None | |
| return b, tg, vv | |
| J.joint_loss = joint_loss_rssm | |
| J.prep_joint = prep_joint_rssm | |
| if __name__ == "__main__": | |
| J.main() | |
Xet Storage Details
- Size:
- 2.66 kB
- Xet hash:
- e73c9e516ce7a94a5f9cb230fdb9a54128d60eec1824d8eddd2de2193297b21b
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.