AI Papers Reader

Personalized digests of latest AI research

View on GitHub

When an AI Agent Rehearses Its Own Environment

Training large language models to act as agents—booking flights, querying orders, calling APIs—usually means paying for executable environments that are hard to build and verify, or leaning on separate simulators that can be inaccurate and hard to ground. A research team from Shanghai Jiao Tong University, Zhejiang University, Tencent and other institutions proposes a different route: let the agent practice inside its own imagination. In a preprint, they introduce EnvACE, an agentic reinforcement-learning method that trains a single policy to both act and role-play the environment’s response.

At each turn, EnvACE’s policy first generates a tool call in an acting role. It then switches to a rehearsal role and produces the response that call would plausibly trigger—an API result, an error message, a user reply—before deciding its next action. The authors call this “world rehearsal.” Acting and rehearsal share the same parameters and are optimized together with a role-wise Group Relative Policy Optimization (GRPO) objective that uses separate reward baselines for each role. In the main experiments, the team trained Qwen3-8B on the CM2 dataset with up to 30 interaction turns.

A case study in the paper illustrates the mechanism. An EnvACE-trained agent is asked to check order W4284542. In one turn, it plans to call get_order_details with that ID, rehearses the likely response, realizes the call would fail because the ID needs a leading #, and revises the argument to #W4284542 before executing. EnvScaler-8B and a vanilla agent, by contrast, make the invalid call, receive “Error: Order not found,” and need extra turns plus user clarification to recover.

Across the BFCL-v4 function-calling benchmark, the multi-domain τ2-Bench, and VitaBench, EnvACE-8B reached an overall average of 32.91%, beating the environment-scaling baselines EnvScaler-8B by 0.99 percentage points and AWM-14B by 0.37 points. On VitaBench, its 16.0% average was the best among 7B–8B methods. On FinMCP-Bench, a financial tool-use benchmark, EnvACE’s tool F1 (TF1) was 46.78%, the best among compared systems, with the highest tool precision at 54.04%. Controlled τ2-Bench experiments showed world rehearsal added 5.5 percentage points over standard GRPO, raising the average from 31.2% to 36.7%; sharing parameters between acting and rehearsal contributed a further 1.2 points over a per-role-policy variant.

The internalized world model can also be used at test time. With two private rehearsal attempts in parallel, EnvACE improved its overall score from 36.7% without rehearsal to 40.9%, whereas using the base model for rehearsal produced no comparable gain. The budget mattered: at three attempts, performance fell, which the authors attribute to longer inputs approaching context limits. They also note that the study evaluated models only up to 8B parameters and focused on tool-interactive tasks.

If these results generalize, the authors argue, world rehearsal could decouple agent training from the cost and brittleness of external environments—letting models make and correct mistakes in a private, imagined world before committing to actions in the real one. That conclusion remains contingent on larger-scale validation, but it points to a training signal that comes from the agent itself.