ChilleD/agent_world_model_env

1---2title: Agent World Model Environment Server 3emoji: 🤖4colorFrom: blue5colorTo: indigo6sdk: docker7pinned: false8app_port: 80009base_path: /web10tags:11  - openenv12---13 14# Agent World Model15 16AgentWorldModel-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.17 18 19## Quick Start20 21You can interact with the AWM environments at Huggingface Space : [ChilleD/agent_world_model_env](https://huggingface.co/spaces/ChilleD/agent_world_model_env) 🤗.22 23### 1. Start the Server24 25```bash26# From the OpenEnv root directory27PYTHONPATH=src:envs uv run uvicorn envs.agent_world_model_env.server.app:app --host 0.0.0.0 --port 889928```29 30### 2. Connect with the Client31 32```python33import asyncio34from agent_world_model_env import AWMEnv35from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction36 37async def main():38    async with AWMEnv(base_url="http://localhost:8899") as env:39        # Reset to a scenario with a specific task40        result = await env.reset(scenario="e_commerce_33", task_idx=0)41        print(f"Task: {result.observation.task}")42        print(f"Tools available: {result.observation.num_tools}")43        print(f"Verifier support: {result.observation.has_verifier}")  # {sql: True, code: True}44 45        # List available tools46        tools = await env.list_tools()47        for tool in tools[:3]:48            print(f"  - {tool.name}: {tool.description}")49 50        # Call a tool51        obs = await env.call_tool("search_products", query="headphones")52        print(f"Result: {obs.tool_result}")53 54        # Run verification (can be called multiple times with different modes)55        result = await env.step(CallToolAction(56            tool_name="verify",57            arguments={"verifier_mode": "code", "final_answer": "optional answer"}58        ))59        print(f"Reward type: {result.observation.reward_type}")60        print(f"Reward: {result.reward}")61        print(f"Verify result: {result.observation.verify_result}")62 63        # End episode (destroys subprocess; set keep_session=True to preserve files)64        result = await env.step(CallToolAction(tool_name="done", arguments={"keep_session": False}))65        print(f"Episode done: {result.done}")66 67asyncio.run(main())68```69 70## Environment Details71 72### Actions73 74AWM supports two action types:75 76| Action | Description |77|--------|-------------|78| `ListToolsAction()` | List all available MCP tools for the current scenario |79| `CallToolAction(tool_name, arguments)` | Call a specific tool with arguments |80 81Special tool names:82- `"verify"` - Run verifier with `{verifier_mode: "sql"|"code", final_answer: "optional"}` arguments83- `"done"` - End the episode and destroy subprocess (does NOT run verifier)84- `"__list_scenarios__"` - List all 1,000 available scenarios and their tasks85 86### Observation Fields87 88| Field | Type | Description |89|-------|------|-------------|90| `reward` | float | Reward value based on reward_type and config |91| `reward_type` | str | Outcome classification (see below) |92| `scenario` | str | Current scenario name |93| `task` | str | Task description in natural language |94| `task_idx` | int | Task index (0-9) |95| `has_verifier` | dict/None | Verifier support: `{sql: bool, code: bool}` or None |96| `num_tools` | int | Number of tools available |97| `tool_name` | str | Name of the tool called |98| `tool_result` | Any | Result from the tool call |99| `error` | str | Error message if any |100| `verify_result` | dict | Verification output after calling verify |101| `trajectory_path` | str | Path to saved trajectory JSON (after `done`) |102| `session_dir` | str | Path to session directory (only if `keep_session=True`) |103 104### Reward Types and Values105 106Default reward configuration:107 108| Type | Reward | Description |109|------|--------|-------------|110| `complete` | 1.0 | Task completed successfully (verifier passed) |111| `incomplete` | 0.1 | Task not completed (verifier failed) |112| `format_error` | -1.0 | Format error (maps from tool_not_found, invalid_args) |113| `tool_not_found` | -1.0 | Tool name not recognized |114| `invalid_args` | -1.0 | Tool arguments invalid |115| Other types | 0.0 | server_error, timeout, etc. |116 117You can customize rewards at reset:118```python119result = await env.reset(120    scenario="e_commerce_33",121    task_idx=0,122    reward_config={"complete": 1.0, "incomplete": 0.0, "format_error": 0.0}123)124```125 126## Session Artifacts127 128When calling `done(keep_session=True)`, the session directory is preserved with:129 130| File | Description |131|------|-------------|132| `trajectory.json` | Full episode trajectory (scenario, task, steps, each action/result) |133| `{scenario}.db` | SQLite database after agent interaction (final state) |134| `{scenario}_initial.db` | SQLite database snapshot before agent interaction |135| `server.py` | Patched Python code for the launched environment |136| `server.log` | Launched environment uvicorn logs (startup + HTTP requests) |137 138When `keep_session=False` (default), all files are cleaned up after the episode.139 140## Verifier Modes141 142AWM supports two verification modes, selected when calling the `verify` tool:143 144### Code Mode (Default, no LLM needed)145 146```python147result = await env.step(CallToolAction(148    tool_name="verify",149    arguments={"verifier_mode": "code", "final_answer": "optional answer"}150))151```152 153Executes a Python verifier function that compares initial and final database states. Deterministic and does not require LLM.154 155### SQL Mode (code-augmented LLM-as-a-Judge)156 157This mode is recommended for judge performance. You need to set the LLM credentials via environment variables before using this mode.158 159```python160# Set LLM credentials via environment variables161# OPENENV_AWM_LLM_BASE_URL, OPENENV_AWM_LLM_API_KEY, OPENENV_AWM_LLM_MODEL162 163result = await env.step(CallToolAction(164    tool_name="verify",165    arguments={"verifier_mode": "sql"}166))167```168 169Runs SQL queries to extract state changes, then uses an LLM judge to determine success.170 171 172## Listing Scenarios & Tasks173 174```python175async with AWMEnv(base_url="http://localhost:8899") as env:176    # List all 1,000 scenarios177    result = await env.step(CallToolAction(tool_name="__list_scenarios__", arguments={}))178 179    print(f"Total scenarios: {result.observation.total}")180    for scenario in result.observation.scenarios[:5]:181        print(f"  - {scenario['name']}: {scenario['num_tasks']} tasks")182        print(f"    Sample task: {scenario['tasks'][0][:80]}...")183```184 185## Server Monitoring186 187The server exposes a `/stats` endpoint for monitoring active sessions:188 189```bash190curl http://localhost:8899/stats191```192 193Returns: `total_sessions`, `max_idle_time_config`, `cleanup_interval_config`, `scenarios` breakdown, and `max_idle_s`.194 195A background cleanup daemon automatically kills sessions idle longer than `MAX_IDLE_TIME` (default 600s) when total sessions exceed `ALLOWED_IDLE_SESSIONS` (default 3000).196 197## Full Agent Interaction Example198 199See [`examples/agent_world_model/example_usage.py`](../../examples/agent_world_model/example_usage.py) for a complete example of an LLM-powered agent that:200 2011. Discovers available tools via `list_tools`2022. Iteratively calls tools to accomplish the task2033. Runs verification via `verify` tool (can use "sql" or "code" mode)2044. Ends episode via `done` action with `keep_session=True` to inspect artifacts205 206The 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.207 208 209## Large-Scale RL Training210 211AWM is designed for large-scale agentic RL. A single server supports thousands of concurrent WebSocket sessions, each with its own isolated environment subprocess.212 213### Simulated Stress Test214 215A stress test simulating large-scale RL is included:216 217```bash218# after server started, then in another terminal:219PYTHONPATH=src:envs uv run python examples/agent_world_model/example_stress_test.py \220    --scale 1024 --concurrency 64 --min-turns 3 --max-turns 20 \221    --think-min 3.0 --think-max 30.0222```223 224This launches 1024 parallel episodes, each with 3-20 multi-turn tool interactions and 3-30s simulated LLM rollout time per turn.225 226## AWM Server Configuration227 228Server configuration is in `server/config.py`, overridable via environment variables:229 230| Config | Default | Env Var | Description |231|--------|---------|---------|-------------|232| `MAX_CONCURRENT_ENVS` | 10000 | — | Max WebSocket sessions |233| `READY_TIMEOUT` | 180s | `OPENENV_AWM_READY_TIMEOUT` | Subprocess startup timeout |234| `MAX_PORT_RETRIES` | 5 | `OPENENV_AWM_MAX_PORT_RETRIES` | Port-retry attempts on startup failure |235| `RETRY_READY_TIMEOUT` | 30s | `OPENENV_AWM_RETRY_READY_TIMEOUT` | Shorter timeout for retry attempts |236| `READY_POLL_INTERVAL` | 0.5s | — | Polling interval during startup check |237| `MAX_IDLE_TIME` | 600s | `OPENENV_AWM_MAX_IDLE_TIME` | Idle session cleanup threshold |238| `ALLOWED_IDLE_SESSIONS` | 3000 | `OPENENV_AWM_ALLOWED_IDLE_SESSIONS` | Session count before idle cleanup triggers |239| `CLEANUP_INTERVAL` | 5s | `OPENENV_AWM_CLEANUP_INTERVAL` | Cleanup daemon scan interval |240 241## Warning242 243AWM treats verifier code and scenario code from the curated [AgentWorldModel-1K](https://huggingface.co/datasets/Snowflake/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.244 245## Citation246 247More details can be found at:248 249| Resource | Link |250|----------|------|251| Hugging Face Space | [ChilleD/agent_world_model_env](https://huggingface.co/spaces/ChilleD/agent_world_model_env) |252| Paper | [arxiv.org/abs/2602.10090](https://arxiv.org/abs/2602.10090) |253| Synthesis Pipeline Code | [Snowflake-Labs/agent-world-model](https://github.com/Snowflake-Labs/agent-world-model) |254| AgentWorldModel-1K | [Snowflake/AgentWorldModel-1K](https://huggingface.co/datasets/Snowflake/AgentWorldModel-1K) |255| Arctic-AWM-4B | [Snowflake/Arctic-AWM-4B](https://huggingface.co/Snowflake/Arctic-AWM-4B) |256| Arctic-AWM-8B | [Snowflake/Arctic-AWM-8B](https://huggingface.co/Snowflake/Arctic-AWM-8B) |257| Arctic-AWM-14B | [Snowflake/Arctic-AWM-14B](https://huggingface.co/Snowflake/Arctic-AWM-14B) |258 259If you find this work useful, please kindly cite:260 261```bibtex262@article{wang2026agentworldmodelinfinity,263      title={Agent World Model: Infinity Synthetic Environments for Agentic Reinforcement Learning},264      author={Zhaoyang Wang and Canwen Xu and Boyi Liu and Yite Wang and Siwei Han and Zhewei Yao and Huaxiu Yao and Yuxiong He},265      year={2026},266      eprint={2602.10090},267      archivePrefix={arXiv},268      primaryClass={cs.AI},269      url={https://arxiv.org/abs/2602.10090},270}271```272