Skip to content

Fix truncated-horizon evaluation masks - #3

Draft
Arison591 wants to merge 1 commit into
shaohua-pan:mainfrom
Arison591:agent/slice-eval-action-pad-mask
Draft

Fix truncated-horizon evaluation masks#3
Arison591 wants to merge 1 commit into
shaohua-pan:mainfrom
Arison591:agent/slice-eval-action-pad-mask

Conversation

@Arison591

Copy link
Copy Markdown

Summary

  • Slice action_is_pad to the requested action_horizon alongside the ground-truth action chunk.
  • Add a focused regression test covering a shorter evaluation horizon with padded actions.

Root cause and impact

evaluate_batch() supports evaluating fewer action steps than are stored in the batch. It already sliced action_gt to that horizon, but passed the full-length padding mask to action_mse(). For example, evaluating a 2-step horizon from a 4-step chunk failed while expanding a [B, 4] mask over a [B, 2, D] tensor:

RuntimeError: The expanded size of the tensor (2) must match the existing size (4) at non-singleton dimension 1

The mask now follows the same horizon as the target, so truncated-horizon metrics both run and continue to exclude padded steps.

Validation

  • python -m pytest -q tests/test_metrics.py1 passed
  • python -m ruff check starwam/training/metrics.py tests/test_metrics.pypassed
  • python -m compileall -q starwam testspassed
  • Repeated the regression on an NVIDIA GeForce RTX 5090 (SM 12.0, PyTorch 2.7.0+cu128) — returned the expected {"action_mse": 1.0}

No model weights, dataset, or external service are required for the regression test.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant