Skip to content

Snowflake-Labs/agent-world-model

Folders and files

NameName
Last commit message
Last commit date

Latest commit

ย 

History

2 Commits
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 

Repository files navigation

AWM Logo Agent World Model

Infinity Synthetic Environments for Agentic Reinforcement Learning

Zhaoyang Wang1, Canwen Xu2, Boyi Liu2, Yite Wang2, Siwei Han1,
Zhewei Yao2, Huaxiu Yao1, Yuxiong He2

1UNC-Chapel Hill ย  2Snowflake AI Research ย 

arXiv HuggingFace HuggingFace HuggingFace HuggingFace

Agent World Model (AWM) is a fully synthetic environment generation pipeline that synthesizes 1,000 executable, SQL database-backed tool-use environments exposed via unified MCP interface for large-scale multi-turn agentic reinforcement learning.


๐Ÿ“ฃ News

  • Mar 16, 2026: we added the verification demo, please refer to Verification section!
  • Feb 10, 2026: we open-sourced the synthesis pipeline, 1,000 synthesized environments and RL trained agents at Huggingface!

๐ŸŽฏ Overview

The AWM synthesis pipeline incldues:

  1. Start from a high-level scenario (e.g., "an online shopping platform")
  2. Generate user tasks that serve as functional requirements
  3. Synthesize a SQLite database (schema + sample data) as the state backend
  4. Generate a Python interface layer (FastAPI + MCP) as the action/observation space
  5. Generate verification code that inspects database state changes for reward signals

๐Ÿ”ฎ Resources

We released the syntheszied 1,000 executable environments and corresponding tasks, databases, and verification in huggingface. Please checkout huggingface repo at Snowflake/AgentWorldModel-1K.

Resource Link
๐Ÿ“„ Paper ๐Ÿ“„ arxiv.org/abs/2602.10090
๐Ÿ’ป 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 want to directly use our synthesized environments, please download by

hf download Snowflake/AgentWorldModel-1K --repo-type dataset --local-dir ./outputs/

Then you can skip to Environment Management and Agent Demo to start using the environments with the agent demo.

๐Ÿ“ฆ Setup

Run uv sync to setup the python environment. And set your LLM API credentials:

# OpenAI or any other compatible services
export AWM_SYN_LLM_PROVIDER="openai"
export OPENAI_API_KEY="your-api-key"
# optional, if you are using a custom base url
export OPENAI_BASE_URL="http://xxxxxx"

# Azure OpenAI
export AWM_SYN_LLM_PROVIDER="azure"
export AZURE_ENDPOINT_URL="https://your-endpoint.openai.azure.com/"
export AZURE_OPENAI_API_KEY="your-api-key"

# configure the model/LLM for synthesis
export AWM_SYN_OVERRIDE_MODEL="your-model-name such as gpt-5"

๐Ÿ”ฅ Synthesis

AWM CLI

All synthesis is exposed through the awm command-line tool. Run awm --help to see available commands:

awm --help

Available commands:
  gen        Synthesis pipeline commands
  โ”œโ”€โ”€ scenario   Generate scenario names from seed set
  โ”œโ”€โ”€ task       Generate user tasks per scenario
  โ”œโ”€โ”€ db         Generate database schema and create SQLite databases
  โ”œโ”€โ”€ sample     Generate and insert sample data into databases
  โ”œโ”€โ”€ spec       Generate API specification for each scenario
  โ”œโ”€โ”€ env        Generate MCP environment code
  โ”œโ”€โ”€ verifier   Generate verification code for tasks
  โ””โ”€โ”€ all        Run the full synthesis pipeline
  env        Environment management commands
  โ”œโ”€โ”€ start      Start MCP server for a scenario
  โ”œโ”€โ”€ check      Check if an MCP server is running and list its tools
  โ”œโ”€โ”€ check_all  Check all generated environments
  โ””โ”€โ”€ reset_db   Reset databases to initial state
  agent      Run a tool-use agent to solve a task by interacting with the environment
  verify     Verify agent run outputs using code-augmented LLM-as-a-Judge or purely code-based Judge

Use awm <command> --help to see options for any command, e.g. awm gen task --help.

Step 1: Scenario Generation

We start with a seed set of scenarios and generate 1,000 unique scenario descriptions. Note that only the names are used as seeds; the descriptions are included in the seed file for ease of use.

export EMBEDDING_OPENAI_API_KEY="your-api-key for the embedding model"

awm gen scenario \
    --input_path outputs/seed_scenario.jsonl \
    --output_path outputs/gen_scenario.jsonl \
    --target_count 1000

Step 2: Task Generation

We generate 10 tasks per scenario, which are also serving as the requirements for building the environment.

awm gen task \
    --input outputs/gen_scenario.jsonl \
    --output outputs/gen_tasks.jsonl

Step 3: Database Synthesis

We define the database schema and complete the initial state to fully support the generated tasks.

# database schema
awm gen db \
    --input outputs/gen_tasks.jsonl \
    --output outputs/gen_db.jsonl

# sample data for initial state
awm gen sample \
    --input_task outputs/gen_tasks.jsonl \
    --input_db outputs/gen_db.jsonl \
    --output outputs/gen_sample.jsonl

Step 4: Interface Synthesis

We first generate API spec for better generating the Python code of the environment in MCP interface.

# API spec (interface schema)
awm gen spec \
    --input_task outputs/gen_tasks.jsonl \
    --input_db outputs/gen_db.jsonl \
    --output outputs/gen_spec.jsonl

# Environment code
awm gen env \
    --input_spec outputs/gen_spec.jsonl \
    --input_db outputs/gen_db.jsonl \
    --output outputs/gen_envs.jsonl

Step 5: Verification Synthesis

We provide two options for verification:

  1. code-augmented LLM-as-a-Judge (sql)
  2. purely code-based Judge (code)
awm gen verifier \
    --mode sql \
    --input_task outputs/gen_tasks.jsonl \
    --output outputs/gen_verifier.jsonl

Environment Management

Run and check each environment. The MCP endpoint will be available at http://localhost:8001/mcp.

# Reset databases to initial state
awm env reset_db \
    --input_db outputs/gen_db.jsonl \
    --input_sample outputs/gen_sample.jsonl

# Start MCP server for a scenario
awm env start \
    --scenario "scenario_name" \
    --envs_load_path outputs/gen_envs.jsonl \
    --port 8001

# Check if MCP server is running
awm env check --url http://localhost:8001/mcp

# Batch test all generated environments
awm env check_all --output outputs/gen_envs.jsonl

Agent Demo

AWM includes a simple agent demo that connects to an MCP environment to solve tasks via multi-turn tool calling. Please start the environment and use vLLM to serve the model before running the agent.

# serve the model
vllm serve Snowflake/Arctic-AWM-4B --host 127.0.0.1 --port 8000

# start the environment, this will create an isolated folder outputs/servers/<timestamp> to save the environment related files such as initial.db, final.db, and etc.
awm env start --scenario e_commerce_33 --envs_load_path outputs/gen_envs.jsonl --port 8001

# run the agent
awm agent \
    --task "show me the top 10 most expensive products" \
    --mcp_url http://localhost:8001/mcp \
    --api_url http://localhost:8000/v1 \
    --model Snowflake/Arctic-AWM-4B

Verification

AWM supports two types of verification:

  1. sql, the recommended code-augmented LLM-as-a-Judge (requires LLM env vars, see Setup section)
  2. code, purely code-based Judge
# launch an isolated environment and run the agent to finish the task of corresponding scenario
awm agent \
    --scenario e_commerce_33 \
    --task_id 0 \
    --api_url http://localhost:8000/v1 \
    --model Snowflake/Arctic-AWM-4B

# after interaction, the trajectory will be saved to outputs/agents/<timestamp>, we can verify it by
awm verify --input outputs/agents/<timestamp> --mode sql

Citation

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

About

Agent World Model: Infinity Synthetic Environments for Agentic Reinforcement Learning

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages