feat: undersample episodes by loss distribution

[naidenovaleksei]

Loss-Based Dataset Resampling

This PR adds infrastructure to resample training data based on per-sample loss scores from a previously trained model, flattening the loss distribution so the model sees harder examples more often.

How it works

1. Score the dataset — run a trained model in inference-only mode using the new policy_loss_only / pretrained_loss_only configs. These write a parquet file of per-sample losses indexed by meta/sample_id.

2. Resample at training time — the training dataset pipeline now has two new steps:

   • left_join_parquet — joins the sample list with the scores parquet on meta/sample_id
   • drop_overrepresented_by_loss — bins samples by mean loss and caps each bin to max_bin_freq × N, discarding overrepresented low-loss samples

What changed

New utilities (src/rmind/utils/pipeline.py)

• left_join_parquet: attaches external parquet columns to a DataFrame by key
• drop_overrepresented_by_loss: caps loss bins to flatten the loss distribution; no-op if scores aren't present

Per-sample loss in objectives (components/objectives/)

• Added LOSS as a new ObjectivePredictionKey computed with reduction="none" (per sample, not aggregated)
• Added reduction_none context manager in base.py to temporarily switch all loss modules
• Implemented per-sample loss output in policy, forward_dynamics, inverse_dynamics, and memory_extraction
• Bug fix in policy.py: score_logprob was computing log-prob against the predicted mean instead of the ground truth

New configs

• config/inference/yaak/control_transformer/policy_loss_only.yaml — run policy model to score dataset
• config/inference/yaak/control_transformer/pretrained_loss_only.yaml — run pretrained model to score dataset
• config/resampling/policy.yaml / pretrained.yaml — which loss columns to use for resampling
• config/datamodule/yaak/predict_train.yaml — predict dataloader over the training split

Dataset pipeline (config/_templates/dataset/yaak/train.yaml, train_debug.yaml)

• Added the two new pipeline steps (join + drop) at the end of sample processing
• clamp_brake_pedal and drop_pedal_conflict_episodes are now opt-out flags (default on) instead of hardcoded
• headings_denoised/heading moved to the flat SELECT output (was only in spatial join args before)

Loss module fixes (components/loss.py)

• LogitBiasFocalLoss / LogitBiasCrossEntropyLoss: register logit_bias as a proper buffer (survives device moves and serialization); handle None bias correctly
• GramAnchoringLoss: refactored to compute per-sample losses and added reduction attribute

Prediction writer (callbacks/prediction/dataframe.py)

• Added final_path parameter to move the consolidated parquet to a specific output location

Usage

1. PT

# predict over train set to get per-sample losses
just predict  inference=yaak/control_transformer/pretrained_loss_only datamodule=yaak/train_debug model.artifact=[pt_baseline_model_artifact] "+scores_output=$HOME/scores_pt_debug.parquet"

#  train with undersampling
just train experiment=yaak/control_transformer/pretrain datamodule=yaak/train_debug  +scores_parquet=$HOME/scores_pt_debug.parquet +loss_max_bin_freq=0.02

2. FT

# predict over train set to get per-sample losses
just predict  inference=yaak/control_transformer/policy_loss_only datamodule=yaak/train_debug model.artifact=[ft_baseline_model_artifact] "+scores_output=$HOME/scores_ft_debug.parquet"

# train with undersampling
just train experiment=yaak/control_transformer/finetune datamodule=yaak/train_debug model.artifact=[tt_encoder_model_artifact] +scores_parquet=$HOME/scores_ft_debug.parquet +loss_max_bin_freq=0.02

[felixmaximilian]

1. *_loss_only config default to datamodule: yaak/predict. But predict.yaml has no meta/sample_id , only datamodule=yaak/train_debug has it. If used as is, will likely fail.
2. (Latent) In-place += in LogitBias* losses. input += self.logit_bias mutates a tensor that, on the LOSS path, is a Rearrange view sharing storage with raw_logits. use input = input + self.logit_bias instead.
3. (Nice-to-have) No unit tests for left_join_parquet / drop_overrepresented_by_loss — pure, correctness-critical, testable with in-memory polars frames.

[felixmaximilian]

Please collect statistics how the new filter affects the number of episodes to ensure we are not accidentally dropping high signal datapoints

[felixmaximilian]

Make the loss modules honor the full reduction contract they're imitating (Gram: handle "sum"; Focal: expose a real reduction instead of hardcoded .mean()), so reduction_none behaves uniformly and there are no silent fall-throughs (applies to GaussianNLLLoss, LogitBiasCrossEntropyLoss)