UniMASK: Unified Inference in Sequential Decision Problems — part 4/4

Cleaned-text code points [47020, 50391). Bibliography omitted here; source.md includes it.

In reward-conditioned evaluation, choosing reasonable return-to-go (RTG) tokens on which to condition is non-trivial: asking for large reward in cases in which the goal is very close to the starting state leads to impossible-to-satisfy queries. Conversely, using the average reward as the goal return-to-go might be too conservative for easy initializations in which the point mass object can traverse most of the maze. To obviate this problem, we try to automatically determine what a reasonable RTG is at evaluation time using the following method: 1) reset the environment (leading to a random initial state); 2) find the trajectory in the dataset which has the most similar initial state (which also includes information about the goal location), and its total reward R R ; 3) Condition on an RTG of 1.1 × R 1.1\times R . We found this behaves as intended qualitatively.

Hyperparameters.

For each model and task, hyperparameters were obtained with a random-search method, which sweeped over batch sizes ( 50,100,200 50,100,200 ), token embedding dimensions ( 64,128 64,128 ), number of layers ( 2 2 , 3 3 , or 4 4 ), number of heads ( 8 8 , 16 16 ), state loss re-scaling factors ( 1 1 , 0.5 0.5 , 0 0 ), and learning rates (selected log uniformly between 10 − 5 10^{-5} and 10 − 3 10^{-3} ). With number of layers, we refer to attention layers for transformers, and hidden layers for feedforward models.

Given that we found little difference between various hyperparameters across model types, the same set of hyperparameters was used across all conditions. The final hyperparameters are as follows: 10 − 4 10^{-4} learning rate, 100 100 batch size, 128 128 embedding dimension, 4 4 layers, 16 16 attention heads, and a state loss re-scaling factor of 1 1 (equivalent to no re-scaling).

Similarly to the Minigrid experiments, each finetune model used the same hyperparameters as its corresponding single-task model, with the learning rate lowered to 8 × 10 − 5 8\times 10^{-5} , and the number of epochs lowered to 600 600 .

Computational cost.

We used the same compute infrastructure for our Maze2D experiments as in the Minigrid experiments (described in Appendix E ). A training run in Maze2D takes approximately 4 hours on our server, but more than 20 runs can be run in parallel. In total, running all runs for Table 1 should take on the order of 10 hours when using our setup in parallel.

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .