Question about potential inconsistency in latent_action selection between training and evaluation

[shiyouheng]

Hello,

I've been looking through the codebase and noticed a difference in how latent_action is extracted during training versus evaluation. I wanted to ask about this to better understand the implementation, as it might be a deliberate design choice that I'm not fully grasping.

In vla-scripts/finetune_libero.py, during training, latent_action_tokens are selected using a mask based on action_gt > 32000:

def action_decoder_forward(self, batch, vla_output):
    # ... (other code)
    action_gt = batch["labels"].to(latent_tokens.device)
    mask = action_gt > 32000

    latent_action_tokens = []
    for idx, per_sample_latent_tokens in enumerate(latent_tokens):
        per_sample_latent_action_tokens = per_sample_latent_tokens[mask[idx], :]
        latent_action_tokens.append(per_sample_latent_action_tokens)
    # ... (uses latent_action_tokens for prediction)

However, in the evaluation code (experiments/robot/libero/run_libero_eval.py), latent_action is selected using the positions of the last four tokens, which is equivalent to using generated_ids > 32000 to create a mask for extracting latent_action:

def predict_latent_action(...):
    # ... (other code)
    latent_action = latent_tokens[:, -4:]
    # ...

Since the LlamaForCausalLMbackbone uses next token prediction, I'm wondering if these two approaches might be extracting latent_action from different positions, creating a potential mismatch between training and evaluation.

I also noticed another section of code that accounts for the next token prediction behavior by using the mask with action_gt, which aligns with my understanding of how this should work:

# Compute Accuracy and L1 Loss for Logging
action_logits = output.logits[:, wrapped_model.module.vla.vision_backbone.featurizer.patch_embed.num_patches : -1]
action_preds = action_logits.argmax(dim=2)
action_gt = batch["labels"][:, 1:].to(action_preds.device)
mask = action_gt > 32000

# Compute Accuracy
correct_preds = (action_preds == action_gt) & mask
action_accuracy = correct_preds.sum().float() / mask.sum().float()

Would you be able to clarify whether the difference in latent_action selection between training and evaluation is a deliberate design choice, a potential oversight, or perhaps something I'm misunderstanding about the code flow?

Thank you very much for your time and help—I really appreciate it!

[retsuh-bqw]

The mismatch arises from the use of history latent actions. During training, history latent action tokens are assigned -100 in batch["labels"] (no loss computation) and are filtered out with the condition action_gt > 32000. However, during evaluation, if history latent actions are still included, filtering based on detokenized indexes will yield 8 tokens (history + current) instead of the expected 4.

[shiyouheng]

Thank you so much for your reply.
I think my earlier question wasn't as precise as it should have been, so let me try to reframe it. What I'm particularly curious about is the potential off-by-one discrepancy in latent action selection between training and evaluation, stemming from the next-token prediction behavior of the LlamaForCausalLM backbone.
During training, when we use action_gt > 32000 to create the mask for extracting latent_action_tokens, are we effectively selecting tokens at position t that correspond to the ground truth at t?
But since the model is trained with next-token prediction (predicting t+1 from t), shouldn't the latent action tokens used for prediction be from position t to predict the action at t+1?
In evaluation, when we take latent_tokens[:, -4:], it seems like we're capturing the most recent tokens that would align with the predicted next step. This makes me wonder if the training selection (action_gt > 32000) and evaluation selection ([-4:]) might be offset by one position due to the autoregressive nature of the model.

[retsuh-bqw]

They retrieve tokens from the same position. The only difference comes from the teacher-forcing mechanism during training, but its impact is negligible in practice.

[feiyu12138]

Hi,
I also have the same concern. Becase: hidden_states[:,t] refers to output_logits[:,t+1]. I'm afraid this may cause indexing mismatch in training/ evaluation. Could author please confirm if there do exists some problem or it's my misunderstanding?

Thanks a lot