Action-conditioned JEPA future predictor (optional action_dim)

[ShauryaVM]

Small wedge toward an action-sensitive world model.

FutureFeatureMapPredictor / WorldActionModel now take an optional action_dim. When set, actions are projected into the same seed space as the history and added as a residual (Delta-JEPA style). The action projection is zero-initialised, so at init the forecast matches the history-only path for any action — same containment idea as the reasoning coupling gate.

Default construction is unchanged (action_dim=None). Opt in via world_model_kwargs={"action_dim": 2} (or whatever control dim you use) and pass actions= into predict_future / forward.

How I tested

• pytest Model/tests/test_world_action_model.py (23 passed), including zero-init no-op and post-open sensitivity checks

Made with Cursor

[riita10069]

Under the contribution standard proposed in #198, I consider this PR AI slop in its current form. It adds an architecture and API without wiring them into the training loop or reporting any trained JEPA result or model-performance evaluation. This is not ready for review; please provide the required real-world validation before resubmitting.

[ShauryaVM]

The action residual is now in the training loop, not just the WAM API.

AutoE2E.forward teacher-forces trajectory_target into predict_future when world_model_kwargs.action_dim is set. train_il / wf_train_il take action_conditioned_jepa (default off) which builds action_dim=128 (flattened 64×(a,κ) plan) and stores it on the checkpoint so eval rebuilds the same architecture.

A/B on a Combined mock batch, 12 AdamW steps, CPU (JSON: Model/evaluation/results/action_jepa_ab.json):

	JEPA	notes
history-only Combined	0.0725	action_dim=None
action-conditioned	0.0688	Δ −0.0037 vs history-only
same model, shuffled actions	0.0688	Δ vs matched ≈ 0 on random frames

Containment holds: at init, matched vs shuffled actions are byte-identical (pred L1 = 0, JEPA 0.0982 both). After those 12 steps action_proj L2 = 2.40 and swapping actions moves the forecast (pred L1 = 0.010). Gradients reach action_proj through jepa_loss.

I am not claiming a KITScenes JEPA win — future frames in this batch are random, so shuffled actions barely change reconstruction. The number that matters for landing the wiring is: the residual is a no-op until trained, then the train loop actually uses the plan.

Same CLI on a packed partition:

python -m evaluation.action_jepa_ab --shard-dir /path/to/partition --out action_jepa_ab.json

How I tested: pytest Model/tests/test_action_jepa_ab.py Model/tests/test_world_action_model.py (4 + 23 passed).