PaperlayerView exact revision on arXiv

Source-linked AI summary

Paying Less Generalization Tax: A Cross-Domain Generalization Study of RL Training for LLM Agents

Zhihan Liu, Lin Guan, Yixin Nie, Kai Zhang, Zhuoqun Hao, Lin Chen, Asli Celikyilmaz, Zhaoran Wang, Na Zhang

arXiv:2601.18217v1cs.AIcs.LG

TL;DR

Generalist agents are trained on narrow environments but must operate across unknown domains, creating a need to understand what preserves out-of-domain performance. The paper analyzes environment and modeling factors, validates state information richness through lightweight observation randomization, and finds that richer states, complex planning, and explicit reasoning support generalization while mid-training can harm uncovered domains.

  • Problem

    The paper asks which RL environment properties and modeling choices preserve agents’ performance in unseen domains when post-training data covers only a narrow task slice.

  • Method

    The study compares four RL environments, analyzes state information richness and planning complexity, intervenes with goal-irrelevant observation content, and examines mid-training and step-by-step reasoning.

  • Results

    State information richness and planning complexity correlate with cross-domain generalization; state randomization improves robustness, while explicit reasoning preserves it and mid-training risks degradation on uncovered domains.

  • Takeaways & Limitations

    When target environments are unknown, prioritize training environments with richer states and complex planning, use lightweight state randomization, and enable explicit reasoning.

  • Takeaways & Limitations

    The evidence uses a limited domain suite, and the initial characterization of effective environments is relatively ad hoc rather than conclusively causal.

Abstract

from arXiv · show

Generalist LLM agents are often post-trained on a narrow set of environments but deployed across far broader, unseen domains. In this work, we investigate the challenge of agentic post-training when the eventual test domains are unknown. Specifically, we analyze which properties of reinforcement learning (RL) environments and modeling choices have the greatest influence on out-of-domain performance. First, we identify two environment axes that strongly correlate with cross-domain generalization: (i) state information richness, i.e., the amount of information for the agent to process from the state, and (ii) planning complexity, estimated via goal reachability and trajectory length under a base policy. Notably, domain realism and text-level similarity are not the primary factors; for instance, the simple grid-world domain Sokoban leads to even stronger generalization in SciWorld than the more realistic ALFWorld. Motivated by these findings, we further show that increasing state information richness alone can already effectively improve cross-domain robustness. We propose a randomization technique, which is low-overhead and broadly applicable: add small amounts of distractive goal-irrelevant features to the state to make it richer without altering the task. Beyond environment-side properties, we also examine several modeling choices: (a) SFT warmup or mid-training helps prevent catastrophic forgetting during RL but undermines generalization to domains that are not included in the mid-training datamix; and (b) turning on step-by-step thinking during RL, while not always improving in-domain performance, plays a crucial role in preserving generalization.

1 Introduction

The paper studies how to preserve cross-domain generalization when agents are post-trained on limited environments but deployed across unknown domains. It identifies environment and modeling choices associated with transfer and evaluates state randomization as a practical intervention.

  • Motivation: Post-training covers limited domains because online RL environments are costly to build, while deployment spans broader and shifting task distributions.This motivates selecting training environments that preserve capabilities in unseen domains.
  • Motivation: The paper distinguishes cross-domain generalization from in-domain generalization and focuses on transfer across substantially different tools, states, or goals.Preliminary experiments show unseen-environment performance can drop as in-domain scores improve during RL.
  • Environment Factors: Across four RL environments, state information richness and planning complexity strongly correlate with cross-domain performance.Richness measures the information the agent must process; planning complexity is estimated through goal reachability and trajectory length under a base policy.
  • State Randomization: The study intervenes on state information richness by injecting small amounts of goal- or domain-irrelevant content into observations without changing the task.The method increases information richness while leaving other task-space components unchanged.
  • Modeling Choices: Mid-training or SFT warmup preserves knowledge in covered domains but can degrade generalization to domains excluded from the mid-training mixture.Explicit step-by-step thinking does not always improve in-domain scores but helps preserve generalization.
  • Implications: The paper recommends training environments with richer states and more complex planning, alongside explicit reasoning and lightweight state randomization, when target domains are unknown.The characterization is based on four environments and is presented as indicative rather than conclusive.

2 Related Work

Related work covers interactive RL frameworks for LLM agents, cross-domain transfer studies, and augmentation methods for robustness. This paper extends these lines by examining which environments best train generalist agents and by applying semantic textual distraction to agentic RL.

  • Interactive Agent RL: Recent work extends RL frameworks to multi-turn LLM-agent settings and introduces collaborative multi-turn benchmarks.These efforts build on broader RL-agent advances, including AlphaGo.
  • Cross-Domain Generalization: Prior studies measure transfer from training domains to unseen domains and compare how SFT and RL preserve cross-domain generalization.One reported comparison finds RL preserves more generalization than SFT on single-turn tasks.
  • Augmentation: State and domain augmentation commonly improve robustness by varying simulator or observation properties, especially in sim-to-real and visual RL.Examples include changes to textures, object properties, lighting, physics, and distractive observations.
  • Augmentation: This work uses state randomization to validate its environmental findings and reports that semantic textual noise with controlled distraction can prevent out-of-domain degradation in agentic RL.The contribution is positioned beyond introducing state randomization itself.

3 Preliminary

The paper formalizes RL training for LLM agents as optimizing expected cumulative reward over language-action trajectories, and distinguishes in-domain validation from cross-domain generalization to unseen environments.

  • RL Training for LLM Agents: An LLM agent observes states and task prompts, generates textual actions, receives rewards, and seeks a policy maximizing expected cumulative trajectory reward.A trajectory contains state, action, and reward tuples across discrete timesteps.
  • RL Training for LLM Agents: Sparse end-of-episode success rewards make direct adoption of single-turn optimization approaches such as GRPO possible for practical agentic tasks.The same framework can extend to multi-turn tasks with sparse rewards.
  • RL Training for LLM Agents: Group-based RL samples multiple full-episode trajectories for each prompt and derives normalized advantages from their reward statistics rather than per-token value functions.GRPO is the cited example of this group-based strategy.
  • RL Training for LLM Agents: GRPO optimizes the policy using an objective with an importance-sampling ratio between the current policy and the sampling policy.The supplied passages identify the ratio but do not provide a readable full objective.
  • Cross-Domain Generalization: In-domain evaluation tests unseen tasks within an environment, whereas cross-domain evaluation tests unseen environments differing substantially in states, actions, or dynamics.The paper also calls cross-domain evaluation out-of-domain (OOD) generalization.
  • Cross-Domain Generalization: Because public agentic domains are limited, experiments train on one environment at a time and evaluate on the remaining environments.This protocol operationalizes cross-domain generalization across available environments.

4 Analytical Framework

The analytical framework evaluates RL generalization across four diverse agentic environments using controlled training and out-of-domain evaluation, with two initial policies supporting model-choice analyses.

  • Environments: The study covers WebShop, Sokoban, ALFWorld, and SciWorld, spanning web navigation, information retrieval, spatial reasoning, and embodied interaction.The environments differ across dimensions including domain realism and action space.
  • Environments: WebShop is a text-based e-commerce simulator requiring page navigation, attribute filtering, and product selection for complex user queries.It represents web navigation and information-retrieval tasks.
  • Environments: Sokoban uses explicit coordinates for walls, players, objects, and targets because the original matrix representation did not yield meaningful RL gains with Llama-8B-instruct.The puzzle requires pushing boxes to targets while avoiding irreversible lock-states.
  • Environments: ALFWorld requires symbolic actions for household interaction, while SciWorld extends it with larger action spaces, specialized tools, and complex multi-stage experiments.SciWorld is included as a significantly more challenging extension of ALFWorld.
  • Evaluation Protocol: RL training runs for 150 steps, checkpoints are saved every 15 steps, and reported success is mean pass@1 over the final four checkpoints across three seeds.The final four checkpoints span steps 105 to 150.
  • Models: Llama-3.1-8B-Instruct is the base model, but near-zero initial success on WebShop and SciWorld requires policy preparation before downstream comparisons.The paper creates Ckpt V1 through WebShop RL and Ckpt V2 through SFT warmup with mixed trajectories.
  • Models: Ckpt V1 and Ckpt V2 serve as downstream initial policies, with Ckpt V2 enabling analysis of SFT warmup or mid-training.WebShop RL raises in-domain success to 34.4% after 20 steps for Ckpt V1.

5 Characterizing Effective Training Domains for Preserving Generalization

RL improves in-domain performance but can reduce performance on unseen domains. Across four environments, higher state information richness and planning complexity align with stronger cross-domain robustness, while domain realism and knowledge absorption do not fully explain the differences.

  • RL training significantly improves in-domain performance but often causes performance drops in out-of-domain environments.
  • The OOD Ranking Score sums each training domain’s performance rank across unseen evaluation domains, with lower scores indicating stronger robustness.Using Ckpt V2, SciWorld ranks first with 3, followed by Sokoban with 5, WebShop with 6, and ALFWorld with 8.
  • State information richness and planning complexity are the two environment properties most aligned with cross-domain generalization.State richness is approximated by average state character count, while planning complexity is reflected by trajectory length and goal reachability.
  • Sokoban and SciWorld combine substantial state information and planning demands, whereas WebShop and ALFWorld are lower in at least one dimension.
  • Sokoban retains stronger OOD performance than ALFWorld and WebShop even when its ∆ID is increased to match or exceed theirs.This indicates that OOD performance is not simply determined by specialized knowledge absorbed during RL.
  • The initial evidence is indicative rather than conclusive because it uses few domains and does not causally isolate the critical environmental factors.

6 From Correlation to Practical Solution for Preserving Generalization

The paper tests whether increasing state information richness can improve generalization without changing task mechanics. Adding controlled, goal-irrelevant information to observations generally improves OOD performance, with reported gains across training domains.

  • State Information Augmentation enriches observations with goal-irrelevant information while preserving task goals, action spaces, and reward functions.
  • The method constructs an augmented state by injecting a controlled volume of irrelevant text into the agent’s input, leaving transitions and rewards unchanged.
  • Insufficient augmentation may have little OOD effect, whereas excessive noise can make learning harder and reduce in-domain performance.
  • In Ckpt V1, augmentation improves OOD performance across all training domains: ALFWorld (+7.0%), WebShop (+35.5%), and Sokoban (+42.5%).
  • In Ckpt V2, augmentation also improves OOD performance for ALFWorld (+7.0%), WebShop (+33.4%), and Sokoban (+5.7%).
  • The results causally validate state information richness as a determinant of agent robustness and support augmentation as a low-cost preservation strategy.

7 Other Influential Factors on Cross-domain Generalization

Modeling choices substantially affect cross-domain generalization: SFT warmup protects knowledge covered by its datamix but increases forgetting elsewhere, while explicit reasoning preserves OOD performance even when it does not improve in-domain scores.

  • SFT Warmup: SFT warmup reshapes retention during subsequent RL, preserving knowledge more robustly in domains covered by its warmup datamix.The comparison uses Ckpt V1 and Ckpt V2 after identical downstream RL training.
  • SFT Warmup: −56.4% versus −11.2%: after Sokoban RL, WebShop success-rate decline was smaller for Ckpt V2 than Ckpt V1.WebShop was included in Ckpt V2’s SFT-warmup datamix.
  • SFT Warmup: +2.4% versus −36.2%: after WebShop RL, Sokoban performance increased for Ckpt V1 but decreased substantially for Ckpt V2.Sokoban was outside Ckpt V2’s SFT-warmup coverage.
  • SFT Warmup: SFT warmup or mid-training therefore trades stronger retention in covered domains for diminished prior capabilities in uncovered domains.The paper recommends broad datamixes and careful control of mid-training when deployment environments are unknown.
  • Step-by-Step Reasoning: Explicit step-by-step reasoning consistently maintains higher OOD performance than reactive action-only baselines.In Ckpt V2 experiments, disabling reasoning caused OOD performance to collapse by over 200% in ALFWorld, Sokoban, and SciWorld training runs.
  • Step-by-Step Reasoning: 1.0% versus 30.5%: when training on ALFWorld and evaluating on WebShop, disabling reasoning sharply reduced success relative to explicit reasoning.The comparison illustrates the OOD benefit of reasoning despite its inconsistent in-domain gains.
  • Step-by-Step Reasoning: Reasoning can leave in-domain performance unchanged or lower, but it reduces reliance on brittle, domain-specific heuristics that do not transfer.On ALFWorld with Ckpt V1, disabling reasoning increased in-domain performance from 73.0% to 77.0%.

8 Limitations and Future Work

The study provides a practical framework for analyzing and improving cross-domain generalization, while identifying broader validation and more formal environment measures as future priorities.

  • Future Work: Future work should test the identified factors across more domains and model scales to assess their universality and uncover additional influences on OOD performance.The authors also propose refining empirical measures with more formal and causally grounded definitions.
  • Future Work: Additional directions include automated environment augmentation, reasoning-structure analysis, adaptive SFT warmup schedules, and theoretical links to generalization and policy optimization.These directions aim to support more principled environment design and scalable, robust agent training.

9 Conclusion

The study finds that state information richness and planning complexity align more closely with out-of-domain performance than realism or surface similarity. State randomization and modeling choices provide practical guidance for preserving capabilities when future test environments are unknown.

  • State information richness and planning complexity, respectively controlling perception and reasoning load, align more closely with out-of-domain performance.
  • The authors causally validate state information richness by injecting goal-irrelevant content into observations without changing the task space.
  • Mid-training or SFT warmup preserves knowledge in covered domains but can worsen generalization on domains excluded from the training mix.
  • Step-by-step thinking during RL and evaluation does not increase in-domain success rates but plays a crucial role in transferring capabilities to unseen environments.
  • Table 8 ranks source domains by success rate on unseen domains, while Tables 9 and 10 report extended-training OOD results for Sokoban and SciWorld.

A.1 RL in ALFWorld

This appendix describes RL implementations and example interactions for ALFWorld, WebShop, Sokoban, and SciWorld. The examples show agents receiving rich textual states, reasoning step by step, and selecting admissible actions.

  • Implementation: RL training and evaluation use verl-agent with eight NVIDIA A100 GPUs for ALFWorld, WebShop, and Sokoban, while SciWorld requires a proxy-client interface.
  • Implementation: ALFWorld and SciWorld allow up to 50 environment steps, whereas WebShop and Sokoban allow up to 15 steps per episode.
  • SFT warmup: SFT warmup uses LlamaFactory and mixes SciWorld, ALFWorld, and WebShop data in a 19:3:1 state-action-pair ratio.
  • Example interactions: ALFWorld states enumerate many room objects and task actions, while WebShop states contain search results, product attributes, and prior observations.
  • Example interactions: The examples require step-by-step reasoning enclosed in <think> tags before an admissible action enclosed in <action> tags.

B.3 Sokoban

This appendix presents Sokoban’s grid-based task representation and example interaction format, alongside state-augmentation implementation details. Sokoban requires constrained box pushing and planning over a coordinate-described grid.

  • Sokoban environment: Sokoban asks agents to push every box onto its corresponding target in a 6×6 coordinate-described grid.
  • Sokoban environment: Agents may move in four directions, cannot pass through walls or pull boxes, and must avoid irreversible deadlocks.
  • Example interaction: The example state lists walls, a goal, a box, a player, and admissible actions, followed by reasoning that selects a move to push the box.
  • Example interaction: SciWorld examples similarly present an embodied task, a detailed observation, admissible actions, and reasoning before activation of a switch.
  • State augmentation: State augmentation implementation includes synthetic distractor objects for ALFWorld and selection of information volume for experiments.

C.4 Information Augmentation Volume Selection

The appendix reports the information-augmentation volume used in Section 6. The referenced table provides the experiment-specific values of ϵ.

  • Table 11 reports the information augmentation volume ϵ used for the experiments in Section 6.
Loading 2601.18217v1…