World Modeling for LLM Agent - Predicting or Acting in the Environment
Agentic reinforcement learning usually trains only on the agent’s own tokens: the commands it writes. The tool outputs it reads, such as query results, error messages and table histories, are masked out. The learning signal is one scalar reward per episode.
ECHO (Shrivastava et al., 2026) and related work (PaW) add a second signal. They put a cross-entropy loss on the observation tokens, so the model also learns to predict what the environment will say. They claim for
- faster learning
- better generalization
- less need for expert demonstrations
The intuition is that an agent that knows what a command will do should choose commands better.
I chose Delta Lake tables on a Databricks-style workspace to test the idea. This is a suitable domain, since most of the small models confuse Delta with plain Parquet and get Unity Catalog names wrong.
This is a little demo of the agent in action:
A quick summary of what I found:
- World model works. On held-out episodes, the ECHO-trained models predict the facts in tool outputs 0.64-0.77 nats per token better than the GRPO-trained one.
- ECHO models solved more tasks. ECHO models beat GRPO by about 0.2, with confidence intervals from +0.12 to +0.29.
- A better world model might not make a better agent.
This project is open source. Check the source at https://github.com/munhouiani/delta_world
Background
I would like to spend some time to connect pre-training to ECHO and its variants.
Pretraining
The model reads a very large amount of ordinary text. At each position it is scored on the probability it gave the token that actually came next: the cross-entropy, $−\log p$. This is where it learns language and facts, such as Japan’s capital.
Supervised Fine-Tuning
The loss is unchanged. Two things change:
- The data: a demonstration, in the chat format the model will be used in, written by a person or a stronger model.
- The mask: the prompt is read but not scored, so the model learns to answer, not to write questions. In an agent’s demonstration the tool outputs are usually masked too. The model learns the agent’s moves, not what the tools said.
Reinforcement Learning
This figure is actually more aligned with GRPO-styled reinforcement learning. The model won’t see any demonstration (unlike supervised-finetuning). The loop has three steps:
- Sample: a sampler runs the current model to answer each prompt several times.
- Score: a verifier gives each answer a reward. GRPO turns rewards into advantages by comparing each answer with its group: above the group’s mean is positive, below is negative.
- Train: the trainer, holding the same model, scores the sampled tokens. Per token, the loss is in effect fine-tuning’s $−\log p$ multiplied by the answer’s advantage $A$, with clipping that limits how far one update moves a probability. A positive $A$ makes the answer more likely, while a negative $A$ makes it less likely. The updated weights go back to the sampler.
A group whose answers all score the same has $A = 0$ and teaches nothing. In agentic RL an episode interleaves the agent’s turns with tool outputs. Only the agent’s tokens are in the loss, and one reward per episode is the only signal.
ECHO
The biggest selling point of ECHO is that ECHO requires the model to also predict the environment. GRPO reads the tool outputs but is never scored on predicting them. ECHO scores them with ordinary cross-entropy, not weighted by the advantage, multiplied by $\lambda$ and divided by the episode’s observation length.
- What it buys: a model that predicts what a command will return has learned something about the environment, a world model, from its own rollouts.
- What it costs: almost nothing. The logits come from the forward pass GRPO already makes, and no extra rollouts are needed.
ECHO Variants
ECHO trains on every observation token, weighted equally. Later work changes which tokens are predicted, how a miss is scored, or how much the loss weighs. Each of the following figures draw a variant from ECHO’s treatment of the tool calls. The lookup gets a fuller output (country: Japan, capital: Tokyo, time: 10:42), and a second call, submit(Tokyo), follows it.
PaW
PaW (Lu et al., 2026) changes all three.
- Which: across an update, it drops the 25% of calls whose actions the policy was surest of, by action entropy. Their actions stay in the RL loss.
- How: it replaces cross-entropy with a clipped mean absolute error. A token costs $1 − p$, never more than 1, and counts only where the model gave it a probability of at most 0.2. They believe a noisy observation can follow the same action in more than one way.
- How much: the weight is one minus the group’s mean return over the maximum, so a group that always succeeds adds nothing.
DGS-WM
DGS-WM (Ran et al., 2026) keeps only what drives the agent’s next action.
- Which: the output is split into minimal units of meaning, such as a sentence. Each unit is scored by how likely the agent’s actual next action is when that unit alone stands in for
the output. Units more than half a standard deviation above the output’s mean are kept. - How much: a kept unit’s cross-entropy is weighted by one minus the model’s confidence on it, so what the model already predicts weighs less.
Qwen-AgentWorld
Instead of adding an auxiliary loss, Qwen-AgentWorld (Qwen team, 2026) trains a separate language world model.
- Format: it trains on over 10 million trajectories. There is no policy and no reward: next-token cross-entropy on the environment’s turns is the whole objective.
- Which: each turn gets a category from statistics such as how much the output overlaps the call. A keep ratio per category decides how much of it is scored: 100% for retrieval, 5% for an echo. Masked turns stay as context.
- Use: the world model then simulates environments for a policy’s RL, or serves as a warm-up stage before it.
This study uses ECHO’s released loss, on every token of every tool output.
Building the Environment
In this project, building an environment costs most of the work, and most of the ways the results could have been wrong. This section describes:
- what the agent acts on
- how a rollout’s outcome is judged
- how models are compared
Simulator
I want my agent learns how to operate on a Databricks workspace: runs serverless SQL and Python over Delta tables in Unity Catalog. However, using a real one as the environment is out of reach due to the following reasons:
- Volume: one GRPO step runs 64 episodes and several hundred tool calls, and a run takes thousands of tool calls.
- Writes and isolation: every episode needs full write access, its own tables, and isolation from the other 63.
- Quota: the free workspace I used for ground truth was driven through a quota-aware runner at 600 calls an hour.
So the agent trains and is evaluated against a local simulator. The real workspace (my personal Databricks Free Edition account) is used only to measure how faithful that simulator is.
The core engine of the simulator is a long-lived Spark Connect server (Spark 4.2.0) with Delta Lake 4.4.1. Each episode gets its own Spark session and schema, and dropped when it ends. The server is never restarted between episodes. Every tables on a Databricks workspace is named as workspace.<schema>.<table>, therefore the engine needs to translates them to a schema unique to the episode. Physical names never appear in an observation, and two rollouts of the same task cannot see each other.
Open-source Spark behaves slightly different than Databricks serverless. Databricks uses Delta as its default table format, and open-source Spark would silently create Parquet tables with no history and no RESTORE. When an operation uses file paths outside Unity Catalog volumes, it should fail as they do on the workspace, with errors captured from it. Databricks uses three-part naming (catalog.schema.table), a direct POSIX-like filesystem paths to structure or unstructured files managed under Unity Catalog volumes (/Volumes), and their cluster and session timezone remains set to Etc/UTC, while open-source Spark has no such restrictions. Therefore, I need a translation layer that makes open-source Spark behaves like Databrick serverless.
How Agent Interact with Environment
Every agent is provided with three tools:
run_sql(query)run_python(code). It runs in a per-episode Python kernel attached to the episode’s Spark session, like a notebook, so temp views and variables carry across cellssubmit(answer). It ends the episode
Tasks
In this project, I designed 5 tasks families, and every task instance is generated from a seed: tables, data, history, damage and prompt.
| Family | The task | What the agent must work out | What test-ood-hard pushes |
|---|---|---|---|
| A, time travel | A bad overwrite, delete or merge damaged a table, and other jobs wrote after it; restore the state just before the damage, with time travel | which commit was the damage, from the history and its metrics | 6-9 commits, the good state 3-5 versions back, 2-4 distractor commits of the damage’s own kind |
| B, MERGE | A nightly merge through a view over a staging table duplicates or drops rows; fix the view (highest seq wins; rows without an id are ignored) |
that staging repeats keys and has null ids, from counts and merge metrics | 3-5 repeated keys and 2-4 null ids in 12-20 staging rows |
| C, schema evolution | A batch’s schema changed (an added, widened or renamed column); load it so a stated downstream read works, changing the schema in place | the enforcement error the naive insert raises, or a rename that succeeds silently by position | larger tables and batches |
| D, clustering | A dashboard query reads every file; fix the layout so it can skip most of them | no clustering, or clustering on the wrong column (which a plain OPTIMIZE does not fix) |
60k-90k rows, and at most 15% of rows scanned instead of 25% |
| E, data quality | Answer a question about a table with duplicates, NULLs and distractor columns | the data itself | 400-800 rows, 20% NULLs, 30-45% dirty values, 5-12% duplicates |
Here are examples from each task family.
A, Time Travel
A faulty MERGE into
workspace.ops.invoicesoverwrote existing rows with bad values and inserted bogus rows. Other jobs have written to the table since then. Restore the table to exactly the state it was in immediately before that operation, using Delta time travel. Do not rebuild or rewrite the table by hand. Call submit when you are done.
The agent is expected to first call DESCRIBE HISTORY (agent sees all columns, only three columns are illustrated here):
1 | version | operation | operationMetrics (abridged) |
The damage is version 5, so the answer is the version before it, not the one before the latest. The agent should then run
1 | RESTORE TABLE workspace.ops.invoices TO VERSION AS OF 4 |
B, MERGE
A nightly job merges the staging table
workspace.crm.subscribers_stagingintoworkspace.crm.subscribersthrough the viewworkspace.crm.subscribers_updates:
MERGE INTO workspace.crm.subscribers t USING workspace.crm.subscribers_updates s ON t.id = s.id WHEN MATCHED THEN UPDATE SET * WHEN NOT MATCHED THEN INSERT *Run as it is, the job fails with an error. Fix the view so that the job is correct for any staging data, under these rules: when an id appears more than once in staging, the row with the highest
seqwins; rows without an id are invalid and must be ignored. Keep the job statement as it is. Then run the job once, check the result, and call submit.
The agent is expected to first check the staging table via a SELECT statement and be able to notice the staging has two kinds of bad rows.
1 | SELECT * FROM workspace.crm.subscribers_staging |
1 | id | name | visits | seq |
Run as it is, the job stops on the repeated 801, which is already in the target: [DELTA_MULTIPLE_SOURCE_ROW_MATCHING_TARGET_ROW_IN_MERGE] Cannot perform Merge as multiple source rows matched and attempted to modify the same.... Therefore, the agent needs to run a fix that keeps the latest row per id and drops the NULL ids:
1 | CREATE OR REPLACE VIEW workspace.crm.subscribers_updates AS |
Finally, the agent will run the MERGE statement in the prompt.
1 | MERGE INTO workspace.crm.subscribers t USING workspace.crm.subscribers_updates s ON t.id = s.id WHEN MATCHED THEN UPDATE SET * WHEN NOT MATCHED THEN INSERT * |
The job then updates 3 rows and inserts 4.
1 | | num_affected_rows | num_updated_rows | num_deleted_rows | num_inserted_rows | |
C, Schema Evolution
A new batch has landed in
workspace.events.sessions_batch, and its schema no longer matchesworkspace.events.sessions. Load every row of the batch intoworkspace.events.sessions, keeping the existing rows, so that this downstream read works:
SELECT id, label, value, device FROM workspace.events.sessions ORDER BY idThe read needs
device(NULL for the existing rows). Change the table’s schema in place; do not drop or recreate it. Call submit when you are done.
The agent is expect to first use DESCRIBE statement to inspect the schema:
1 | DESCRIBE DETAIL workspace.events.sessions |
The agent should notice the batch has one more column than the table, so the obvious load fails:
1 | INSERT INTO workspace.events.sessions SELECT * FROM workspace.events.sessions_batch |
1 | [DELTA_METADATA_MISMATCH] A metadata mismatch was detected when writing to the Delta table. |
A passing solution alters the table first, then loads the batch:
1 | ALTER TABLE workspace.events.sessions ADD COLUMNS (device STRING); |
D, Clustering
This dashboard query reads far more of
workspace.ads.clicksthan it should, and it gets slower every week:
SELECT count(*) AS n, round(sum(amount), 2) AS total FROM workspace.ads.clicks WHERE campaign_id = 344Fix the table’s layout so that the query can skip most of the table’s files. Keep the table and its rows as they are: do not drop, replace or rewrite it by hand, and do not delete rows. Call submit when you are done.
In this example, the table holds 39,000 rows in 5 files. DESCRIBE DETAIL shows no clustering columns, and the query reads every row. A plain OPTIMIZE changes nothing: it adds no files. A passing solution clusters on the filtered column and rewrites the layout:
1 | ALTER TABLE workspace.ads.clicks CLUSTER BY (campaign_id); |
The table then has 15 files, and the query reads 6.7% of its rows, under the 25% limit.
E, Data Quality
Answer a question about the table
workspace.warehouse.inventory. Some rows repeat an earlier row exactly in every column exceptid. Drop those repeats (treating NULLs as equal to each other), then give the total ofscoreover the remaining rows. Ignore NULLs. Investigate the table with queries, then call submit with the answer as a single value.
In this example, the table has 211 rows. Its channel column holds three values in 32 spellings, plus 20 NULLs:
1 | id | grade | score | channel | region |
The question says exactly, so the spellings must be left as they are here. A passing answer, submitted as 9194:
1 | SELECT coalesce(sum(score), 0) |
Verifiers
Reward for each task family is binary. When an episode end, the reward is calculated by a verifier runs in the agent’s own session. The agent never sees the verifier. The verifier inspect the the final state of the table, but never parse what the agent says it did. Family E’s submitted answer is the one exception. This is crucial as the agent might claim “restored via time travel” without actually running the RESTORE statement.
| Family | Passes only if |
|---|---|
| A | the contents equal the snapshot of the last good version, the table is the same table (same id), and every commit since setup is a RESTORE or maintenance |
| B | tonight’s result is right; and then, with held-out staging data swapped in and the target rewound, re-running the job through the agent’s view gives the right result again |
| C | the same table, no replace or overwrite commit, the expected columns and types, and the downstream read returns the expected rows |
| D | the same table and contents, the filtered column among the clustering keys, only layout commits since setup, and emulated min/max data skipping reads at most 25% of rows (the fix reaches 5-11%) |
| E | the answer matches: numbers within 1e-6, strings exactly, ignoring case and quotes |
Training Setup
Model and Training Framework
- Policy: NVIDIA Nemotron 3 Nano 4B, non-thinking. It is a hybrid that carries most of its state in Mamba-2 layers, with four attention layers.
- RL: NeMo RL 0.7 (Megatron backend), GRPO with 8 prompts × 8 rollouts per step, up to 16 turns, 1,024 new tokens per turn and an 8,192-token context.
- Sampling: temperature 0.6 and top-p 1.0, in training and evaluation alike.
- The sampler’s engine: vLLM, eager, with prefix caching off. With CUDA graphs on, it corrupted the hybrid model’s generation.
- Length: 90 steps, with every arm seeing the same prompts at each step (seed 42). Checkpoints every 10 steps.
- Hardware: each arm used one H200 and its own 32-CPU episode service.
- Cost: a 90-step run took about 4 hours and cost about $23 on Modal.
Training Recipe
I follow ECHO’s released configuration (microsoft/echo-rl).
- no KL penalty to the starting policy;
- the policy loss averaged per sequence, not per token;
- ratio clipping 0.2 on both sides.
Three Arms Ran
To test the idea, I trained 3 models:
- GRPO
- ECHO with $\lambda=0.005$
- ECHO with $\lambda=0.02$
$\lambda$ is the weight for observation loss. In the paper, they use $\lambda=0.02$, and I ran an additional arm with $\lambda=0.005$, which is a quarter of it.
What Agent Reads And What Loss Trains On
Every tool result is rendered into text by one function, used identically in training and evaluation.
- Format: tables print as
col | colrows, in a deterministic order unless the query sorts them. - Masking: timestamps, durations, ids and storage paths become
<TS>,<MS>,<ID>and<PATH>. - Truncation: outputs over 800 tokens keep their head and tail around
... N rows omitted .... - Errors: they keep their class and message and lose their stack traces.
Let’s further discuss what tokens in a tool result are predictable, and whats are not.
- Trivial: boilerplate, and anything derivable from context. Column headers, repeated rows, and echoes of the agent’s own command or of the task text all count:
Overwritehere repeats the task’s “INSERT OVERWRITE”. - Unpredictable: masked values such as
<TS>. - Predictable: the rest. These are the facts the action and the table state determine, such as operation names, write modes, row and file counts, byte counts and error classes. Version numbers are trivial, since each one is the previous plus one.
During training, these three kinds of tokens are included in observation loss. However, the following sections will report predictable tokens separately. This is because a model can lower its loss on the whole output by copying headers and repeated rows from context (reward hacking). The predictable tokens are the ones a world model has to get from the action and the state.
ECHO’s loss, as released, trains on every token of every tool output’s body, except for two things: the chat template’s scaffold around each tool message, and harness messages such as the error for a malformed tool call.
Observation and Findings
ECHO Learns World Model
I took 178 dev-split episodes, which are never seen during training, feed into model checkpoints, and then measured each model’s cross-entropy on the observation tokens and predictable tokens.
| Predictable tokens | All observation tokens | |
|---|---|---|
| GRPO | 1.89 | 1.10 |
| ECHO $\lambda=0.005$ | 1.25 | 0.58 |
| ECHO $\lambda=0.02$ | 1.13 | 0.45 |
The tables shows both ECHO arms are far below GRPO on predictable tokens.
To see how the improvement looks like, token by token, I took one row of DESCRIBE HISTORY output from a held-out episode, and feed it into model checkpoints. Each predictable token is shaded by how surprised each model was to see it.
| Token | GRPO | ECHO $\lambda=0.02$ |
|---|---|---|
operation WRITE |
7.5 | 0.8 |
mode Append |
7.6 | 0.03 |
partitionBy: [] |
4.8 | 0.2 |
numFiles: 6 |
3.7 | 2.9 |
numOutputBytes: 508344 (6 tokens, mean) |
2.3 | 2.3 |
numOutputRows: 37000 (5 tokens, mean) |
2.0 | 2.7 |
Values are surprise in nats.
- The structural facts become near-certain. A table loaded once by an append has a
WRITEinAppendmode with no partitioning. - The exact counts stay uncertain for every model, as they should: the byte count cannot be derived from anything the agent has seen. ECHO is no better on them, and slightly worse on the row count here.
The previous probe shows how surprised model is by an output it is shown. Next, I would like to know a model can predict a failure. The next problem will let a model reads the start of a scripted demonstration and writes the next output itself, so it has to commit to what the command will do, including whether it fails.
The clearest example is from family-C. Recall that tasks from Family C are related to schema change. The task says the batch’s schema no longer matches the table, and the model has seen the schemas of both tables:
1 | > DESCRIBE TABLE workspace.iot.readings > DESCRIBE TABLE workspace.iot.readings_batch |
and the batch rows (808 | alpha | 8410650166, …). Three facts in those outputs decide what an insert does here:
- The
workspace.iot.readingstable storesvalueasint. - The
workspace.iot.readings_batchtable holdsvalueasbigint. - The
workspace.iot.readings_batch‘s values are too big for anINT, which stops at 2,147,483,647.
The next command, an INSERT statement, looks ordinary, so a model has to connect it to what it read.
1 | INSERT INTO workspace.iot.readings SELECT * FROM workspace.iot.readings_batch |
and the actual result is
1 | [CAST_OVERFLOW_IN_TABLE_INSERT] Fail to assign a value of "BIGINT" type to the |
All context will be given to both GRPO and ECHO models. The GRPO model outputs
1 | id | kind | value |
which means GRPO predicts the insertion will success. While ECHO model outputs
1 | [CHANGED_COLUMN_NAME_OR_TYPE] Cannot append column `value` to table |
ECHO model gets the outcome, but names a different error class: the real error is an overflow raised by the values, not a type check. Its message uses facts 1 and 2, with no sign of fact 3. I ran multiple probes, and ECHO model predicts the failure every time, but GRPO never does.
ECHO Models Have More Tasks Solved, Mostly Through Two Habits GRPO Abandoned
Each point in the figure is a model checkpoint. Moving left means a better world model (less surprised by the observation tokens). Moving up means more tasks solved (pass@1). We can see that both ECHO models are better world models and solving more tasks.
Above figure shows pass@1 in totals, which hide the arms differ. Let’s see the result from another perspective:
We now can clearly see that GRPO and ECHO models perform almost equally good in most task families, except for family C and E. Upon checking the GRPO model at different steps, these situations are caused by GRPO forgetting its learned behaviour.
Let’s start with Family E. In short, GRPO stopped calling submit. E’s answer travels in submit. A-D are judged on the table’s state, and a turn with no tool call ends the episode, so for them submit is optional. In GRPO’s training batches:
| Step | Rollouts that call submit | Rollouts ending a turn with ““ and no call | Reward 1 |
|---|---|---|---|
| 5 | 29/64 | 0/64 | 22/64 |
| 30 | 1/64 | 48/64 | 43/64 |
| 40 | 16/64 | 39/64 | 40/64 |
At step 30, GRPO ended its episodes like this one, a family-A task it solved:
1 | "... which restored exactly the 10 rows deleted by the errant version 8 operation. ... |
</think> closes a reasoning block, out of place with thinking off. On A-D this ending cost nothing; on E it cost every reward.
- GRPO had been learning E. Its E rollouts won 13 of 24 at step 5 and 15 of 24 at steps 15 and 20. From step 30 they won none.
- The ECHO arms learned it. They submitted more as training went on: 51 and 36 of 64 rollouts at step 30, and 0.8-0.95 later. Neither ended a turn on
</think>.
For family C, GRPO stopped altering the schema first. C’s batch adds, widens or renames a column. The reliable order is to change the table’s schema, then insert. On dev at step 90:
| GRPO | ECHO $\lambda=0.005$ | ECHO $\lambda=0.02$ | |
|---|---|---|---|
| ALTER before the first insert | 0/96 | 80/96 | 86/96 |
| First insert fails | 64 | 9 | 12 |
| … and no insert succeeds afterwards | 49 | 8 | 7 |
Here is a typical GRPO failure, and ECHO $\lambda=0.02$ on the same instance:
1 | Task: A new batch has landed in workspace.events.sessions_batch, and its schema no longer |
- GRPO reads the error, fixes its cause, and never retries. Its stated reasoning has degenerated into copied filler (“I’ll keep the existing rows.”), and it reports a load that did not happen.
- ECHO fumbles the syntax three times, but in the right order.
- The insert’s form is not the cause. ECHO $\lambda= 0.005$ names its target columns in every solved C episode, but ECHO $\lambda= 0.02$ inserts by position, as GRPO does: 86 of 96 episodes, 74 of them solved. GRPO’s positional inserts solve 33 of 94.
- GRPO abandoned altering first habit. GRPO’s training batches still did altering first in 9 of 16 C rollouts at step 30, and in none of 64 at steps 60 and 90. The ECHO arms kept it and grew it, ending at 83-90% of evaluation episodes.
A Better World Model Doesn’t Stop Inventing Commands
Both ECHO arms write DESCRIBE statements that do not exist, and GRPO never does:
Episodes with a failing DESCRIBE |
A | B | C | D | E |
|---|---|---|---|---|---|
| GRPO | 0 | 0 | 0 | 0 | 0 |
| ECHO $\lambda= 0.005$ | 39 | 25 | 19 | 9 | 1 |
| ECHO $\lambda= 0.02$ | 0 | 2 | 66 | 77 | 3 |
ECHO $\lambda= 0.02$ opens 110 of its 480 dev episodes with DESCRIBE CLUSTERING OF <table>, which the simulator rejects as a syntax error, on C as well as D. On C it moves on after one error. On D it sometimes cannot: 9 of its 96 D episodes repeat one failing call five times or more, and 44% reach the turn limit, against none for GRPO. Example taken from one dev episode:
1 | Task: This dashboard query reads far more of workspace.support.tickets than it should ... |
Knowing what a parse error looks like did not stop it from producing one, again and again. ECHO $\lambda= 0.02$ ‘s D score still ties GRPO’s (0.45 against 0.46): when it gets past the opening, it clusters as well.
Less Tool Usage Is Not Guaranteed
| Arm | Tool calls per episode | Turns when solved | Tool-error rate | Turn limit reached | No submit | Output tokens per episode |
|---|---|---|---|---|---|---|
| GRPO | 8.0 | 6.5 | 0.23 | 0.30 | 0.30 | 717 |
| ECHO $\lambda= 0.005$ | 4.6 | 4.9 | 0.14 | 0.03 | 0.04 | 558 |
| ECHO $\lambda= 0.02$ | 5.5 | 4.4 | 0.22 | 0.17 | 0.17 | 397 |
Prior work expects the observation loss to cut turns, tool errors and tokens. Here the ECHO arms do use fewer, but family by family the difference is GRPO’s loops on C and E. On D, ECHO $\lambda= 0.02$ uses more than twice GRPO’s calls (9.5 against 4.2).
Learning Speed Is Indistinguishable
One of the method’s main claims is faster learning.
| Paired with GRPO, dev | Step 30 | Step 60 | Step 90 |
|---|---|---|---|
| ECHO $\lambda= 0.005$ | +0.21 | +0.10 | +0.20 |
| … without E | +0.06 | −0.06 | +0.02 |
| ECHO $\lambda= 0.02$ | +0.01 | +0.17 | +0.22 |
| … without E | −0.09 | +0.02 | +0.10 |
- On the families that do not depend on
submit, neither arm learned faster. - Family A reversed for ECHO $\lambda= 0.005$ at step 60 (−0.29), and stayed behind.
Conclusion: Why The World Model May Not Be What Helped
The paper reports gains on terminal tasks with Qwen3 8B and 14B, and my ECHO arms gained too. That alone doesn’t show that a world model makes a better agent, for three reasons:
- Most of our gain looks like regularization: the observation loss stopped GRPO’s collapse.
- Our Delta Lake tasks lack some conditions the paper’s result depends on. Copying its recipe exactly wouldn’t change that.
- Our test also differs from the paper’s: far fewer rollouts, a higher learning rate and a smaller model.
The released recipe has no KL penalty to the starting policy, and without one GRPO collapsed onto short outputs. GRPO model’s cross-entropy on tool output ended above its base model, entropy halved and outputs shrank to a third. It also dropped two habits its base model had: calling submit, which its base model does in 46% of dev episodes, and altering the schema before inserting, which its base model does in 66% of family C episodes.
The observation loss works against this. It trains the model to predict what the environment wrote, on every step and on the same weights the policy uses, which the policy loss doesn’t do.
Also, the observation loss doesn’t stop the model from learning. The ECHO arms submit and alter first more often than their base models do. The loss stopped the collapse, and that left RL free to strengthen the habits that win.
The way to find out is to train GRPO with a KL penalty, on this recipe, and compare it with the same ECHO arms. If it closes the gap, the effect is regularization. Any gap that remains could be the world model.