Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning

Thank you for your review! We answer your questions below.

Q1. Types of environments.

R1: We use a set of MuJoCo environments that is standard in the field of meta-RL and consistent with our baselines. To further demonstrate the generality of BATI, we have additionally conducted a preliminary multi-agent experiment. We choose Kuhn Poker, a two-player card game with discrete state and action spaces, differing from the continuous MuJoCo environments used in our paper. We generate different player-2 (opponent) policies as "tasks" and learn an adaptive policy for player-1 over 10 episodes (20 steps) of contexts. As shown in the table below, BATI maintains superior performance over all baselines, further showcasing its capabilities and generalization.

Q2. Comparison with single-model, transformer-based approaches.

R2: We note that Algorithm Distillation imposes additional requirements on the training data (i.e. generated by the learning histories of an RL algorithm), making it less general than BATI and unsuitable for our evaluation setup. In general, we remark that the family of single-model approaches performs Bayesian inference to deduce the posterior given a context trajectory, and is thus similarly flawed as the context encoders of our baselines, as argued in Sec. 3 and 4 in our paper.

To support this argument, we add a new, recent baseline DPT that also makes use of a transformer to encode the context and predict the optimal action, outperforming Algorithm Distillation. Due to the computational requirements of training transformers and the tight rebuttal schedule, we implement DPT and evaluate it in AntDir and HalfCheetahVel, as shown below. It can be seen that DPT is outperformed in both environments.

Q3. Design and capacity of task embeddings.

R3: We demonstrate the expressitivity of our task embedding design with extensive experiments in our paper, showing that BATI outperforms the baselines across the board. To further support the capabilities of BATI, we conduct a new OOD task generalization experiment in AntDir, which is even more demanding for the task embedding, requiring generalization to both OOD contexts and tasks. We sample training goal directions from $[0, \pi)$ and testing directions from $[\pi, 2\pi)$ to build disjoint distributions of training and testing tasks. The results are shown in the table below. It can be seen that while the performance decreases compared to the in-distribution case, BATI still outperforms baselines by a large margin and is the only method to achieve positive returns. This further validates the generalization capability of BATI and its task embedding designs.

Thank you for your efforts in reviewing our paper and the detailed review! We address your concerns below. Due to space limits, all the figures and tables referred below are available anonymously on this website.

Q1. Experiment protocol of BATI and baselines.

R1: We clarify that all methods, including CSRO, are evaluated on the same context distribution $p_\text{test}(M, \mu)$ during test time, as described in Section 5.1. We design our testing protocol to be as challenging as possible since this aligns better with real-world scenarios where dynamics are noisy and the policy may not have control over the context (e.g. provided by a human operator). Other baselines like UNICORN and FOCAL also use a similar protocol. To address your concerns more thoroughly, we evaluate CSRO on HalfCheetahDir using their self-collection protocol and obtain a performance of $-16.2 \pm 25.3$. While the stronger protocol improves performance, CSRO still lags behind BATI by a lot.

Q2. Choice of base offline RL algorithm and hyperparameter tuning.

R2: We switched from BRAC to IQL because we observed in some cases that BRAC simply failed to learn; see website Fig. 2 for results in WalkerRandParams. We also performed little hyperparameter tuning and directly adopted hyperparameters from previous works. The IQL- and CSRO-related ones are borrowed from their respective papers (CSRO CLUB weights are divided by 10 since their implementation has a weight of 10 for the FOCAL loss while ours has 1) while the rest comes from UNICORN, including those of data-collection runs (with learning steps and dataset sizes slightly adjusted to accelerate training).

Furthermore, we would like to remark that as argued in Sec. 3 and 4, the core problem with baselines lies in task inference and not policy learning; the performances of baselines may even worsen with hyperparameter tuning, since the meta-policy would then get better at executing a wrong policy.

Q3. OOD task experiments.

R3: OOD generalization for tasks is extremely challenging, and to the best of our knowledge, the primary experiments of UNICORN and CSRO focus on generalization to OOD contexts on in-distribution tasks, same as BATI. Only UNICORN attempted an OOD task setting with a model-based RL approach, which is out of scope for our paper.

Nevertheless, we acknowledge the importance of OOD generalization and conduct the experiment you requested in AntDir with training goal directions sampled from $[0, \pi)$ and testing directions from $[\pi, 2\pi)$. The results are shown on the website (Tab. 1). It can be seen that while the performance decreases compared to the in-distribution case, BATI still outperforms baselines by a large margin and is the only method to achieve positive returns. This further validates the generalization capability of BATI and its task embedding designs.

Q4. Task embedding table size; sensitivity to hyperparameters.

R4: The size of the task embedding table is the same as the number of training tasks, as described in Section 4.2. To demonstrate the robustness of BATI with respect to the number of training tasks, we conduct an ablation in AntDir that splits the 40 tasks into different train/eval splits. The results are shown on the website (Tab. 2). It can be seen that BATI outperforms baselines in all settings and is highly stable.

We conduct another ablation in WalkerRandParams to validate the robustness of BATI with respect to task embedding size. Compared with the main result of $565.2 \pm 7.9$ with embedding size 32, size 16 and 64 yield performances of $556.4 \pm 1.7$ and $536.2 \pm 6.4$, respectively, demonstrating the robustness of BATI.

Q5. Discussion and comparison with transformer-based methods.

R5: Thank you for the additional related works! We remark that all of the three methods belong to the same paradigm where a transformer is trained to predict actions conditioned on the context, an act of Bayesian inference similarly flawed as the context encoders of our baselines.

To support this argument, we add a new baseline DPT that also uses a transformer to encode the context and predict the optimal action, outperforming AD. Due to limited time, we implement DPT and evaluate it in AntDir and HalfCheetahVel, as shown on the website (Tab. 3). DPT is inferior to BATI in both environments due to spurious correlations.

Q6. Dynamics model training objective $\mathcal{L}\_\text{recon}$.

R6: Thank you for pointing this out! We apologize for the confusion, the correct objective for Eq. 7 should be $\mathcal{L}\_\text{recon}^{X, Z}(\phi, \psi) := \frac{(X^t - g\_\psi(X^b, Z))^2}{\exp h\_\psi(X^b, Z)} + h\_\psi(X^b, Z)$. We design $p(X^t \mid X^b, Z)$ to be a Gaussian distribution with mean and (log) variance parameterized by neural networks g,h, and L_recon is its negative log likelihood (up to a constant). This will be fixed in the revision.

Thank you for your thoughtful and constructive review. We're encouraged by the overall positive assessment that our paper "is generally well-written" and "addresses a critical challenge in ICRL". Below, we address the main concerns raised regarding the generalizability of our method beyond MuJoCo environments and the theoretical connections between our analysis and the proposed BATI approach. As demonstrated in our preliminary multi-agent experiments and theoretical explanations, BATI effectively addresses the identified limitations of previous methods while maintaining strong performance across diverse domains. The following are our detailed responses:

Q1. Results in other domains.

R1: Due to the tight rebuttal schedule, we have conducted a preliminary multi-agent experiment to demonstrate the general applicability of BATI in other domains. We choose Kuhn Poker, a two-player card game with discrete state and action spaces, differing from the continuous MuJoCo environments used in our paper. We generate different player-2 (opponent) policies as "tasks" and learn an adaptive policy for player-1 over 10 episodes (20 steps) of contexts. As shown in the table below, BATI maintains superior performance over all baselines, further showcasing its capabilities and generalization.

Q2. Theoretical analysis of our method.

R2: We apologize for any potential confusion. We establish the preliminaries and analyze the prior methods in Section 3. Building on these insights, Section 4.1 is primarily dedicated to our method that circumvents the prior failure modes. In Section 4.1, we provide two perspectives on the reason why BATI works: 1) the graphical model perspective (Fig. 2) where BATI achieves robustness by blocking $X^b$; 2) the robust likelihood perspective (page 4 right, line 179-219) where the core objective of BATI can be understood as a robust version of the full likelihood that does not depend on $\mu$ through derivations. We will communicate these points more clearly in the revision and greatly appreciate any suggestions to further improve the writing.

Thank you for your kind words and insightful review! We are happy to answer your questions below.

Q1. Training contexts generated by adaptive policies.

R1: The result would likely depend on the exact behavior of the adaptive policy. For example, if the adaptive policy is bad, it may randomly walk around the state space in any task and reveal little information about the true task identity, in which case all approaches would fail due to the uninformative nature of the context. At the other end of the spectrum, if the adaptive policy is good and finds the optimal policy quickly, its behavior may diverge significantly between tasks, resembling the single-task expert case where BATI works and baselines don't. In general, we expect BATI to outperform or at least remain competitive with the baselines for all cases, and the baselines to work only when the adaptive policy behaves similarly in all tasks yet still reveals a lot of task-related information.

Q2. Challenges brought by extremely large task spaces.

R2: A good question! Indeed, this is one of the potential future works we'd like to explore. The task inference procedure of BATI has a time complexity of $\mathcal{O}(NL)$ where $N$ is the number of samples (lower-bounded by the number of training tasks) and $L$ is the length of the context and can be fully vectorized. This allows the inference procedure to run reasonably fast. However, if the task space is extremely large over which we couldn't practically enumerate, more sophisticated techniques could be employed to perform the inference. For example, we can use diffusion posterior sampling to sample from the conditional posterior $p(x \mid y, c)$ where $x$ is the task latent, $y$ is $X^t$, and $c$ is $X^b$. Under this formulation, the task prior $p(x)$ is represented by a diffusion model instead of a table of embeddings, and (robust) likelihood $p(y \mid x, c)$ is provided by the dynamics model whose gradient $\nabla\_x p(y \mid x, c)$ serves as a form of classifier guidance. Leveraging the powerful expressitivity of diffusion sampling procedures, we can efficiently and robustly sample from a much larger task space. However, to demonstrate the advantage of such an approach, a benchmark with a sufficiently complicated task space is required, which is still emerging, as pointed out in your review, and we will explore this further in our future work.

Q3. Connection with BOReL (Dorfman et al., 2022).

R3. Thank you for bringing our attention to this related work! We note that the training data of BOReL are actually augmented via their reward relabelling (RR) procedure, in effect rolling out all context collection policies in all tasks and making $M$ and $\mu$ independent. This avoids the spurious correlation issue during training, but requires additional properties (e.g. tasks differing only by rewards) and ground-truth reward functions. Indeed, as noted in BOReL, "With reward relabelling ablated ... the agent believes the reward is at the point it first visited on the semi-circle" (Page 8, lower left), which is a consequence of the spurious correlation phenomenon. Our BATI can be applied to BOReL to replace RR without any additional requirements.

Q4. Oracle references and number of seeds.

R4: We have updated the main results to include oracle references by single-task experts, available anonymously here (Fig. 1 and Tab. 4). BATI is able to match or even slightly outperform the expert data in most cases due to optimizations performed by IQL. Due to limited time and computational resources, we have run two additional seeds for the main experiments, bringing the total to 5 seeds; the changes are also reflected in the figure and table above.

Q5. Usage of "in-context RL".

R5.: Thank you for the detailed and insightful comments! In BATI, we treat ICRL as mostly synonymous with context-based meta-RL, i.e. learning policies that can efficiently adapt to the context and improve on-the-fly without laborious fine-tuning, regardless of transformer or implicitness. We argue that this "functional" definition better captures the essence of ICRL. For example, there are state-space model approaches that produce a fixed-length representation of the context without transformers, and methods like DPT that are trained to perform explicit Bayesian inference of optimal actions. However, we stress that we greatly value inputs from the community about terminological choices, and are open to changes if deemed necessary.

Q6. Additional hyperparameter and implementation details.

R6: Thank you for your advice! As the ICML 2025 reviewing policy prevents authors from changing the PDF during rebuttal, we'll add relevant details to the appendix for a more self-contained reading experience when we get to update the paper.