Return-to-go conditioning. Unlike previous work that considers return-to-go conditioning [1]}, [2]}, we don't provide the model with many return-to-go tokens (one for each timestep). Providing just the first token should be sufficient for the model to interpret the return-to-go request (as, if necessary, the model can compute the remaining return to go in later steps). We found that this did not negatively impact performance, either for single-task reward-conditioned training, or for randomized masking training.
