ChilleD/agent_world_model_env

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:

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

Reward Types and Values

Default reward configuration:

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:

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:

1. 1.Discovers available tools via list_tools
2. 2.Iteratively calls tools to accomplish the task
3. 3.Runs verification via verify tool (can use "sql" or "code" mode)
4. 4.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:

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:

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},
}