Agent World Model#

AgentWorldModel-1K is a synthetic agentic environment suite containing 1,000 tool-use environments with 10,000 tasks for large-scale RL training. Each environment is a fully functional MCP server with tools, database state, and verification logic.

Quick Start#

You can interact with the AWM environments at Huggingface Space : ChilleD/agent_world_model_env 🤗.

1. Start the Server#

# From the OpenEnv root directory
PYTHONPATH=src:envs uv run uvicorn envs.agent_world_model_env.server.app:app --host 0.0.0.0 --port 8899

2. Connect with the Client#

import asyncio
from agent_world_model_env import AWMEnv
from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction

async def main():
    async with AWMEnv(base_url="http://localhost:8899") as env:
        # Reset to a scenario with a specific task
        result = await env.reset(scenario="e_commerce_33", task_idx=0)
        print(f"Task: {result.observation.task}")
        print(f"Tools available: {result.observation.num_tools}")
        print(f"Verifier support: {result.observation.has_verifier}")  # {sql: True, code: True}

        # List available tools
        tools = await env.list_tools()
        for tool in tools[:3]:
            print(f"  - {tool.name}: {tool.description}")

        # Call a tool
        obs = await env.call_tool("search_products", query="headphones")
        print(f"Result: {obs.tool_result}")

        # Run verification (can be called multiple times with different modes)
        result = await env.step(CallToolAction(
            tool_name="verify",
            arguments={"verifier_mode": "code", "final_answer": "optional answer"}
        ))
        print(f"Reward type: {result.observation.reward_type}")
        print(f"Reward: {result.reward}")
        print(f"Verify result: {result.observation.verify_result}")

        # End episode (destroys subprocess; set keep_session=True to preserve files)
        result = await env.step(CallToolAction(tool_name="done", arguments={"keep_session": False}))
        print(f"Episode done: {result.done}")

asyncio.run(main())

Environment Details#

Actions#

AWM supports two action types:

Action	Description
`ListToolsAction()`	List all available MCP tools for the current scenario
`CallToolAction(tool_name, arguments)`	Call a specific tool with arguments

Special tool names:

"verify" - Run verifier with {verifier_mode: "sql"|"code", final_answer: "optional"} arguments
"done" - End the episode and destroy subprocess (does NOT run verifier)
"__list_scenarios__" - List all 1,000 available scenarios and their tasks

Observation Fields#

Field	Type	Description
`reward`	float	Reward value based on reward_type and config
`reward_type`	str	Outcome classification (see below)
`scenario`	str	Current scenario name
`task`	str	Task description in natural language
`task_idx`	int	Task index (0-9)
`has_verifier`	dict/None	Verifier support: `{sql: bool, code: bool}` or None
`num_tools`	int	Number of tools available
`tool_name`	str	Name of the tool called
`tool_result`	Any	Result from the tool call
`error`	str	Error message if any
`verify_result`	dict	Verification output after calling verify
`trajectory_path`	str	Path to saved trajectory JSON (after `done`)
`session_dir`	str	Path to session directory (only if `keep_session=True`)

Reward Types and Values#

Default reward configuration:

Type	Reward	Description
`complete`	1.0	Task completed successfully (verifier passed)
`incomplete`	0.1	Task not completed (verifier failed)
`format_error`	-1.0	Format error (maps from tool_not_found, invalid_args)
`tool_not_found`	-1.0	Tool name not recognized
`invalid_args`	-1.0	Tool arguments invalid
Other types	0.0	server_error, timeout, etc.

You can customize rewards at reset:

result = await env.reset(
    scenario="e_commerce_33",
    task_idx=0,
    reward_config={"complete": 1.0, "incomplete": 0.0, "format_error": 0.0}
)

Session Artifacts#

When calling done(keep_session=True), the session directory is preserved with:

File	Description
`trajectory.json`	Full episode trajectory (scenario, task, steps, each action/result)
`{scenario}.db`	SQLite database after agent interaction (final state)
`{scenario}_initial.db`	SQLite database snapshot before agent interaction
`server.py`	Patched Python code for the launched environment
`server.log`	Launched environment uvicorn logs (startup + HTTP requests)

When keep_session=False (default), all files are cleaned up after the episode.

Verifier Modes#

AWM supports two verification modes, selected when calling the verify tool:

Code Mode (Default, no LLM needed)#

result = await env.step(CallToolAction(
    tool_name="verify",
    arguments={"verifier_mode": "code", "final_answer": "optional answer"}
))

Executes a Python verifier function that compares initial and final database states. Deterministic and does not require LLM.

SQL Mode (code-augmented LLM-as-a-Judge)#

This mode is recommended for judge performance. You need to set the LLM credentials via environment variables before using this mode.

# Set LLM credentials via environment variables
# OPENENV_AWM_LLM_BASE_URL, OPENENV_AWM_LLM_API_KEY, OPENENV_AWM_LLM_MODEL

result = await env.step(CallToolAction(
    tool_name="verify",
    arguments={"verifier_mode": "sql"}
))

Runs SQL queries to extract state changes, then uses an LLM judge to determine success.

Listing Scenarios & Tasks#

async with AWMEnv(base_url="http://localhost:8899") as env:
    # List all 1,000 scenarios
    result = await env.step(CallToolAction(tool_name="__list_scenarios__", arguments={}))

    print(f"Total scenarios: {result.observation.total}")
    for scenario in result.observation.scenarios[:5]:
        print(f"  - {scenario['name']}: {scenario['num_tasks']} tasks")
        print(f"    Sample task: {scenario['tasks'][0][:80]}...")

Server Monitoring#

The server exposes a /stats endpoint for monitoring active sessions:

curl http://localhost:8899/stats

Returns: total_sessions, max_idle_time_config, cleanup_interval_config, scenarios breakdown, and max_idle_s.

A background cleanup daemon automatically kills sessions idle longer than MAX_IDLE_TIME (default 600s) when total sessions exceed ALLOWED_IDLE_SESSIONS (default 3000).

Full Agent Interaction Example#

See examples/agent_world_model/example_usage.py for a complete example of an LLM-powered agent that:

Discovers available tools via list_tools
Iteratively calls tools to accomplish the task
Runs verification via verify tool (can use “sql” or “code” mode)
Ends episode via done action with keep_session=True to inspect artifacts

The example supports both a local server and the public Hugging Face Space, set AWM_BASE_URL=https://chilled-agent-world-model-env.hf.space (may be slow) to try without local setup.

Large-Scale RL Training#

AWM is designed for large-scale agentic RL. A single server supports thousands of concurrent WebSocket sessions, each with its own isolated environment subprocess.

Simulated Stress Test#

A stress test simulating large-scale RL is included:

# after server started, then in another terminal:
PYTHONPATH=src:envs uv run python examples/agent_world_model/example_stress_test.py \
    --scale 1024 --concurrency 64 --min-turns 3 --max-turns 20 \
    --think-min 3.0 --think-max 30.0

This launches 1024 parallel episodes, each with 3-20 multi-turn tool interactions and 3-30s simulated LLM rollout time per turn.

AWM Server Configuration#

Server configuration is in server/config.py, overridable via environment variables:

Config	Default	Env Var	Description
`MAX_CONCURRENT_ENVS`	10000	—	Max WebSocket sessions
`READY_TIMEOUT`	180s	`OPENENV_AWM_READY_TIMEOUT`	Subprocess startup timeout
`MAX_PORT_RETRIES`	5	`OPENENV_AWM_MAX_PORT_RETRIES`	Port-retry attempts on startup failure
`RETRY_READY_TIMEOUT`	30s	`OPENENV_AWM_RETRY_READY_TIMEOUT`	Shorter timeout for retry attempts
`READY_POLL_INTERVAL`	0.5s	—	Polling interval during startup check
`MAX_IDLE_TIME`	600s	`OPENENV_AWM_MAX_IDLE_TIME`	Idle session cleanup threshold
`ALLOWED_IDLE_SESSIONS`	3000	`OPENENV_AWM_ALLOWED_IDLE_SESSIONS`	Session count before idle cleanup triggers
`CLEANUP_INTERVAL`	5s	`OPENENV_AWM_CLEANUP_INTERVAL`	Cleanup daemon scan interval

Warning#

AWM treats verifier code and scenario code from the curated AgentWorldModel-1K dataset as trusted. Verifier code (server/_verifier_runner.py) is run in a subprocess sandbox (rlimits, restricted builtins, import allowlist); scenario subprocesses run without per-process sandboxing and rely on the container as the outer isolation boundary. The codes are synthetically generated and carefully curated, however, there is no guarantee of absolute safety. We recommend only academic research use.

Citation#

More details can be found at:

Resource	Link
Hugging Face Space	ChilleD/agent_world_model_env
Paper	arxiv.org/abs/2602.10090
Synthesis Pipeline Code	Snowflake-Labs/agent-world-model
AgentWorldModel-1K	Snowflake/AgentWorldModel-1K
Arctic-AWM-4B	Snowflake/Arctic-AWM-4B
Arctic-AWM-8B	Snowflake/Arctic-AWM-8B
Arctic-AWM-14B	Snowflake/Arctic-AWM-14B

If you find this work useful, please kindly cite:

@article{wang2026agentworldmodelinfinity,
      title={Agent World Model: Infinity Synthetic Environments for Agentic Reinforcement Learning},
      author={Zhaoyang Wang and Canwen Xu and Boyi Liu and Yite Wang and Siwei Han and Zhewei Yao and Huaxiu Yao and Yuxiong He},
      year={2026},
      eprint={2602.10090},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2602.10090},
}