Research map / Reinforcement learning / Offline RL
Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes BATI, a model-based maximum-likelihood task inference method for offline in-context RL that infers task latents from dynamics only (not the behavior of the context-collection policy), improving robustness to context shift on MuJoCo.
Problem
Offline in-context RL policies rely on context trajectories and degrade under distribution shift between training and test contexts, because learned encoders approximate Bayesian posteriors that depend on the training distribution of context-collection policies.
Approach
Replace the learned encoder posterior with a maximum-likelihood task inference that maximizes log p(Xt | Xb, M) using a learned table of per-task embeddings and a dynamics estimator; at test time, search over training task embeddings for the one that best reconstructs next-state/reward from the context, conditioning on state-action pairs to block the behavior policy's influence.
Key results
- BATI outperforms UNICORN, CSRO, FOCAL, Recon and DPT across five MuJoCo environments under context shift, converging faster and more stably.
- On AntDir with OOD tasks, BATI achieves 12.1 ± 0.5 episodic return while all baselines are negative (CSRO -45.8, FOCAL -74.7, Recon -51.8, UNICORN -104.6).
- On HalfCheetahDir with noise 0.9 and C=50, BATI reaches -4.7 ± 1.2, rising to 4.7 ± 1.5 at C=400, while baselines stay around -12 to -26.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use maximum-likelihood task inference from dynamics rather than a learned encoder (ablation) | BATI · HalfCheetahDir, OOD contexts and tasks | episodic return: −51.8 ± 37.0 -> 12.1 ± 0.5 | Helped |
| Increase context length from C=50 to C=200 (ablation) | BATI · HalfCheetahDir, ϵ = 0.9 | episodic return: −4.7 ± 1.2 -> 1.8 ± 0.5 | Helped |
| Increase context length from C=200 to C=400 (ablation) | BATI · HalfCheetahDir, ϵ = 0.9 | episodic return: 1.8 ± 0.5 -> 4.7 ± 1.5 | Helped |
| Increase dynamics noise from ϵ = 0 to ϵ = 0.2 (ablation) | learned-encoder baselines · AntDir | episodic return | Hurt |
| Increase dynamics noise from ϵ = 0.2 to ϵ = 0.4 (ablation) | learned-encoder baselines · AntDir | episodic return | Hurt |
| Increase dynamics noise across tested levels (ablation) | BATI · AntDir | episodic return | No clear effect |
| Use the learned Recon encoder instead of BATI's maximum-likelihood inference (ablation) | dynamics reconstruction method · MuJoCo evaluation environments | episodic return | Hurt |
| Evaluate HopperParam with ϵ = 0 instead of noisy dynamics (ablation) | offline ICRL baselines · HopperParam | episodic return | No clear effect |
| Increase the ratio of OOD test contexts in interpolated contexts (ablation) | Recon · HalfCheetahVel | episodic return | Hurt |
| Increase the ratio of OOD test contexts in interpolated contexts (ablation) | BATI · HalfCheetahVel | episodic return | No clear effect |
| Test on tasks with directions from [π, 2π) rather than training directions from [0, π) (ablation) | BATI · AntDir, OOD tasks and contexts | episodic return | Hurt |
Surprising findings
- Baselines' out-of-distribution performance degrades over training in AntDir, indicating learning of spurious behavioral correlations.
- Even without dynamics noise, baselines fail in HopperParam because the dynamics are complex enough that the true correlation is hard to estimate.
Limitations
- Evaluated only on state-based MuJoCo environments; not vision-based or real-world settings.
- Approximation that the integral over behavior policies is task-independent is assumed rather than verified.
- Requires a fixed set of training tasks with per-task embeddings; inference searches over those training embeddings, so OOD tasks are handled only partially.
Experimental setup
Datasets: MuJoCo offline datasets (AntDir, HalfCheetahVel, HalfCheetahDir, HopperParam, WalkerParam) generated with per-task polic
Benchmarks: AntDir HalfCheetahVel HalfCheetahDir HopperParam WalkerParam
Metrics: episodic return
Baselines: UNICORN CSRO FOCAL Recon DPT Oracle
Builds on: IQL UNICORN CSRO FOCAL DPT
Closest papers by meaning
- Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds (NeurIPS 2025 · similarity 0.89)
- Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning (NeurIPS 2024 · similarity 0.89)
- Context Value Informed In-Context Reinforcement Learning (NeurIPS 2026 · similarity 0.89)
- Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.89)
- Entropy Regularized Task Representation Learning for Offline Meta-Reinforcement Learning (AAAI 2025 · similarity 0.89)
- CERTAIN: Context Uncertainty-aware One-Shot Adaptation for Context-based Offline Meta Reinforcement Learning (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.