Environment Modules¶

This section covers how to create and work with environment modules.

Creating Custom Modules¶

Inherit from EnvBase¶

To create a custom environment module, inherit from EnvBase and implement the required methods:

Required Methods to Implement¶

When creating a custom environment module, you must implement:

async def step(self, tick: int, t: datetime) -> None

Execute one simulation step. This is called automatically during AgentSociety simulations. Use this to update time-dependent state.

Parameters:
tick: Duration of this step in seconds t: Current simulation datetime after this step

Tools decorated with @tool

Use the @tool decorator to expose methods as callable functions for agents.

from agentsociety2.env import EnvBase, tool

class MyModule(EnvBase):
    def __init__(self):
        super().__init__()
        # Your initialization

    @tool(readonly=True, kind="observe")
    def get_value(self, agent_id: int) -> str:
        """Get a value for an agent."""
        return f"Value for agent {agent_id}"

    @tool(readonly=False)
    def set_value(self, agent_id: int, value: int) -> str:
        """Set a value for an agent."""
        self._values[agent_id] = value
        return f"Set value for agent {agent_id}"

Reference Implementation¶

For complete reference implementations, see:

SimpleSocialSpace - Social interaction module
PublicGoodsEnv - Public goods game module
PrisonersDilemmaEnv - Prisoner’s dilemma module

The @tool Decorator¶

The @tool decorator marks methods as callable by agents:

from agentsociety2.env import tool

@tool(readonly=True, kind="observe")
def get_status(self, agent_id: int) -> str:
    """Get the status of an agent."""
    return f"Agent {agent_id} status"

Parameters:

readonly (bool): Whether the tool modifies the environment * True = read-only, can be used in queries * False = modifies state, can be used in interventions
kind (str): The type of tool * "observe": Single-parameter observations (requires readonly=True) * "statistics": Aggregate queries (no parameters, requires readonly=True) * None or omitted: Regular tool (any signature, any readonly value)

Tool Types¶

observe: Single-parameter observations (readonly=True required)
statistics: Aggregate queries (no parameters, readonly=True required)
regular: Any other tool (can be readonly or read-write)

See Core Concepts for more details on tool types.

State Persistence and Resume¶

If a module has dynamic in-memory state, such as counters, inboxes, or queues, it can implement workspace persistence so the simulation can continue after an interruption with --resume; see Command Line Interface. The base class provides two overridable hooks, but does not prescribe the state format, file name, or directory. Each module can choose its own layout; state is usually written to <workspace_root>/state/ENV_STATE.json.

Method	Role
`async def to_workspace(self, workspace_path=None)`	Write dynamic state into the workspace. Society calls this through the router at the end of each step. Use `agentsociety2.storage.workspace_state.atomic_write_text` for atomic writes.
`async def restore(self, workspace_path) -> bool`	Restore dynamic state from the workspace on the resume path. The return value indicates whether a snapshot was loaded, allowing the router to distinguish fresh startup from resume. This method is called only after `__init__` and `init()`, so it can safely override resets performed in `init()`.
`def _bind_workspace(self, workspace_path)`	Bind the module to a workspace directory idempotently and create that directory. Called automatically by `to_workspace` / `restore`.
`classmethod async def from_workspace(cls, workspace_path, **init_kwargs)`	Rebuild a ready module by running `cls(**init_kwargs)` and then calling `restore`. During resume, the actor reads `init_kwargs` from `SOCIETY.json`.

Implementation example; compare SimpleSocialSpace, EconomySpace, and MobilitySpace:

import json
import asyncio
from agentsociety2.env import EnvBase, tool
from agentsociety2.storage.workspace_state import atomic_write_text

_STATE_REL = "state/ENV_STATE.json"

class MyPersistentEnv(EnvBase):
    def __init__(self, capacity: int = 100):
        super().__init__()
        self.capacity = capacity          # 构造配置（在 env_kwargs，不存盘）
        self._resources: dict[int, float] = {}  # 动态态（存盘）
        self._lock = asyncio.Lock()       # 不存盘（由 __init__ 重建）

    @tool(readonly=False)
    def allocate(self, agent_id: int, amount: float) -> str:
        """Allocate resources."""
        self._resources[agent_id] = amount
        return f"Allocated {amount} to agent {agent_id}"

    async def to_workspace(self, workspace_path=None) -> None:
        if workspace_path is not None:
            self._bind_workspace(workspace_path)
        atomic_write_text(
            self._workspace_root / _STATE_REL,
            json.dumps(
                {"resources": {str(k): v for k, v in self._resources.items()}},
                ensure_ascii=False,
                indent=2,
                default=str,
            ),
        )

    async def restore(self, workspace_path) -> bool:
        self._bind_workspace(workspace_path)
        state_path = self._workspace_root / _STATE_REL
        if not state_path.is_file():
            return False
        d = json.loads(state_path.read_text(encoding="utf-8"))
        self._resources = {int(k): v for k, v in d.get("resources", {}).items()}
        return True

Conventions:

JSON keys must be strings. For dict[int, ...], write keys with str(k) and read them back with int(k), or use agentsociety2.env.base.dump_int_map / load_int_map.
For pydantic models, use model_dump(mode="json") and model_validate. For datetime values, use isoformat / fromisoformat. Store set values as sorted(list).
Do not persist objects that cannot or need not be serialized, such as asyncio.Lock, external subprocesses, or pretrained model weights. Recreate them in __init__ / init(), as with the routing server in MobilitySpace or the recommender in SocialMediaSpace.
Persistence and replay tables (_agent_state_columns and _env_state_columns) are independent. Replay is append-only time-series data for analysis; workspace persistence is for recovery and overwrites the latest state. See Storage and Result Access.

Registering Modules¶

from agentsociety2.env import CodeGenRouter

router = CodeGenRouter()
router.register_module(MyModule(), name="my_module")

# Or with default name (class name)
router.register_module(MyModule())

Complete Example¶

from typing import Dict
from datetime import datetime
from agentsociety2.env import EnvBase, tool, CodeGenRouter
from agentsociety2 import PersonAgent
from agentsociety2.society import AgentSociety

class WeatherEnvironment(EnvBase):
    """A simple weather environment module."""

    def __init__(self):
        super().__init__()
        self._weather = "sunny"
        self._temperature = 25
        self._agent_locations: Dict[int, str] = {}

    @tool(readonly=True, kind="observe")
    def get_weather(self, agent_id: int) -> str:
        """Get the current weather for an agent's location."""
        location = self._agent_locations.get(agent_id, "unknown")
        return f"The weather in {location} is {self._weather} with {self._temperature}°C."

    @tool(readonly=False)
    def change_weather(self, weather: str, temperature: int) -> str:
        """Change the weather conditions."""
        self._weather = weather
        self._temperature = temperature
        return f"Weather changed to {weather} at {temperature}°C."

    async def step(self, tick: int, t: datetime) -> None:
        """Update environment state for one simulation step."""
        self.t = t
        # Update time-dependent state here if needed

# Use the custom module（agent 以 spec 声明，不实例化）
env_router = CodeGenRouter(env_modules=[WeatherEnvironment()])

agent_specs = [{"id": 1, "profile": {"name": "Bob"}, "config": {}}]

society = AgentSociety(
    agent_specs=agent_specs,
    agent_class_name="PersonAgent",
    env_router=env_router,
    start_t=datetime.now(),
    run_dir=__import__("pathlib").Path("run"),
)
await society.init()

Note

In production, environment routing runs in a dedicated Ray actor (EnvRouterProxy), and all agents share the same environment state. CodeGenRouter can also be used directly in-process. In addition to the default CodeGenRouter, optional routers include ReActRouter, PlanExecuteRouter, TwoTierReActRouter, TwoTierPlanExecuteRouter, and SearchToolRouter (see Architecture and Scalability).

Examples¶

See the built-in environment classes in agentsociety2.contrib.env (class names follow the current code):

SimpleSocialSpace — social interaction
PublicGoodsEnv — public goods game
PrisonersDilemmaEnv — prisoner’s dilemma
TrustGameEnv — trust game