chat_rl: the sampled assistant_end token is masked out, so stopping is never reinforced

In RL training the model never gets a gradient for its decision to stop, because the <|assistant_end|> token it actually sampled is dropped and then re-inserted as padding.

Engine.generate_batch (nanochat/engine.py:288-294) marks a row completed on the terminal token but does not append the token or its mask:

if not completed[i]:
    if token == assistant_end or token == bos:
        completed[i] = True
    else:
        results[i].append(token)
        masks[i].append(mask)

get_batch in scripts/chat_rl.py:130-131 then pads every short row back out using that same token id, with mask 0:

padded_generated_token_sequences = [seq + [assistant_end] * (max_length - len(seq)) for seq in generated_token_sequences]
padded_masks = [mask + [0] * (max_length - len(mask)) for mask in masks]

and mask 0 becomes the ignore index a few lines later (chat_rl.py:138):

targets[mask_ids[:, 1:] == 0] = -1

So the real, sampled stop token is now indistinguishable from padding: same id, mask 0, target -1, and it contributes nothing to pg_obj. Deciding to stop is the one action the reward most directly depends on, and it is the one action that is never reinforced or penalized.

This triggers on any batch where rows finish at different lengths, which for GSM8K rollouts is nearly every batch.

I confirmed it by re-running the two code blocks above standalone on a fake two row rollout: the row that terminated early ends up with mask 0 at the position of its genuine assistant_end, identical to the padded positions after it.

A fix would need generate_batch to keep the terminal token distinguishable from padding, for example by appending it with its true mask before marking the row completed, so downstream code can tell "the model chose to stop here" apart from "this row was padded". That does change what generate_batch returns, and its docstring currently promises the opposite ("Terminal tokens (assistant_end, bos) are not included in the results"), so it may be better handled in chat_rl.py instead. I did not want to guess at which side you would prefer.

Happy to send a PR if you want it fixed in a particular direction.