Creating custom components

Components define the interactive dynamics of an environment. They expand agent action spaces, inject observations, and implement the rules that govern what happens when an agent acts. Any mechanic that does not belong to the scenario’s passive world logic—movement, trading, building, taxation—belongs in a component. Create a custom component when you need agents to perform a new type of action or when you want to add side-channel observations and metrics without modifying an existing scenario.

How components fit into the environment

When a BaseEnvironment subclass is instantiated, each component in the components list is constructed and wired into the environment. At every timestep the environment calls each component in order:

component_step() — execute action effects
generate_observations() — collect per-component observations
generate_masks() — produce action validity masks

At reset time, additional_reset_steps() is called so stateful components can reinitialise their internal bookkeeping.

The Gather component: a worked example

Before writing your own component, read through the Gather component (ai_economist/foundation/components/move.py). It is the canonical reference implementation.

move.py

@component_registry.add
class Gather(BaseComponent):
    name = "Gather"
    required_entities = ["Coin", "House", "Labor"]
    agent_subclasses = ["BasicMobileAgent"]

    def __init__(
        self,
        *base_component_args,
        move_labor=1.0,
        collect_labor=1.0,
        skill_dist="none",
        **base_component_kwargs
    ):
        super().__init__(*base_component_args, **base_component_kwargs)
        self.move_labor = float(move_labor)
        self.collect_labor = float(collect_labor)
        self.skill_dist = skill_dist.lower()
        self.gathers = []  # stateful: reset in additional_reset_steps

Step-by-step: implementing a custom component

Subclass BaseComponent and set class attributes

Every component must declare four class-level attributes before __init__ is called:

Attribute	Type	Purpose
`name`	`str`	Unique identifier used by the registry. Must not contain `.`.
`component_type`	`str` or `None`	Optional shorthand (e.g. `"Trading"`). Used when looking up via `env.get_component()`.
`agent_subclasses`	`list[str]`	Agent class names this component interacts with. Must be non-empty.
`required_entities`	`list[str]`	Resources, landmarks, or endogenous variables this component expects to exist in the world.

from ai_economist.foundation.base.base_component import (
    BaseComponent,
    component_registry,
)

@component_registry.add
class BonusOutput(BaseComponent):
    name = "BonusOutput"
    component_type = "Production"
    agent_subclasses = ["BasicMobileAgent"]
    required_entities = ["Coin", "Labor"]

agent_subclasses must list concrete agent names (e.g. "BasicMobileAgent", "BasicPlanner"), not base class names. If you list more than one, none of them may be a subclass of another.

Implement __init__

Accept your custom kwargs before passing *base_component_args and **base_component_kwargs to super().__init__().After super().__init__() you can access:

self.world — the World object (maps + agents + planner)
self.n_agents — number of mobile agents
self.resources / self.landmarks — lists of entity names
self.episode_length — total timesteps per episode
self.inv_scale — shared inventory scaling factor for observations

def __init__(self, *base_component_args, bonus_rate=0.1, **base_component_kwargs):
    super().__init__(*base_component_args, **base_component_kwargs)
    self.bonus_rate = float(bonus_rate)
    assert 0.0 <= self.bonus_rate <= 1.0
    self._bonus_log = []  # internal tracker, reset in additional_reset_steps

Implement get_n_actions

Tell the environment how many actions (excluding NO-OP) this component adds for each agent type.

def get_n_actions(self, agent_cls_name):
    """
    Return the number of actions (not including NO-OPs) for agents of
    type agent_cls_name.

    Returns:
        None   -- agent does not participate
        int    -- single action space with that many choices
        list   -- multiple action spaces as [("name", n), ...]
    """
    if agent_cls_name == "BasicMobileAgent":
        # One binary action: trigger bonus output or not
        return 1
    return None

For a planner that sets per-agent tax levels across n agents with k brackets:

def get_n_actions(self, agent_cls_name):
    if agent_cls_name == "BasicPlanner":
        return [("Tax_{}".format(i), 10) for i in range(self.n_agents)]
    return None

Implement get_additional_state_fields

Declare any extra keys that should live in agent.state and their reset values. The environment calls this during construction and again at every reset.

def get_additional_state_fields(self, agent_cls_name):
    """
    Return {state_field: reset_value} for agents of type agent_cls_name.
    Fields are merged into agent.state and reset to reset_value each episode.
    """
    if agent_cls_name == "BasicMobileAgent":
        return {"bonus_output_count": 0}
    return {}

Return {} for agent types that do not need extra state. Never return None.

Implement component_step

Apply the effects of each agent’s chosen action. This is called once per timestep, after all component steps in the list before yours have run.

def component_step(self):
    world = self.world
    step_log = []

    for agent in world.agents:
        action = agent.get_component_action(self.name)
        if action == 0:  # NO-OP
            continue
        if action == 1:
            bonus = int(np.random.rand() < self.bonus_rate)
            agent.state["inventory"]["Coin"] += bonus
            agent.state["endogenous"]["Labor"] += 1.0
            agent.state["bonus_output_count"] += bonus
            if bonus:
                step_log.append({"agent": agent.idx, "bonus": bonus})

    self._bonus_log.append(step_log)

Access the agent’s chosen action with agent.get_component_action(self.name). Action index 0 is always NO-OP.

Implement generate_observations

Return a dict of {agent_idx: obs_dict}. Only include agents for which this component provides observations. The structure must stay consistent across timesteps.

def generate_observations(self):
    """
    Returns:
        obs (dict): {agent.idx: observation_dict}
    """
    obs = {}
    for agent in self.world.agents:
        obs[str(agent.idx)] = {
            "bonus_output_count": agent.state["bonus_output_count"],
            "bonus_rate": self.bonus_rate,
        }
    return obs

Return string keys (str(agent.idx)) so that the environment’s observation collation works correctly.

Implement generate_masks

Return a binary mask dict indicating which actions are valid. The mask length must equal the number of non-NO-OP actions. NO-OP is always valid and must not appear in the mask.The base class provides a sensible default (all actions valid). Override only when you need to restrict actions based on state.

def generate_masks(self, completions=0):
    """
    Returns:
        masks (dict): {agent.idx: np.ndarray of shape (n_actions,)}
    """
    masks = {}
    for agent in self.world.agents:
        # Disable bonus action if agent already used it this step
        can_act = 1.0 if agent.state["bonus_output_count"] == 0 else 0.0
        masks[agent.idx] = np.array([can_act], dtype=np.float32)
    return masks

For multi-action-space planners, return a nested dict:

# masks["p"] = {"Tax_0": np.ones(10), "Tax_1": np.ones(10), ...}

Implement additional_reset_steps (optional)

Reset any internal state that persists across timesteps. Called automatically during env.reset().Looking at Gather for reference:

def additional_reset_steps(self):
    # From Gather in move.py
    for agent in self.world.agents:
        if self.skill_dist == "none":
            bonus_rate = 0.0
        elif self.skill_dist == "pareto":
            bonus_rate = np.minimum(2, np.random.pareto(3)) / 2
        elif self.skill_dist == "lognormal":
            bonus_rate = np.minimum(2, np.random.lognormal(-2.022, 0.938)) / 2
        agent.state["bonus_gather_prob"] = float(bonus_rate)
    self.gathers = []

For BonusOutput:

def additional_reset_steps(self):
    self._bonus_log = []

Implement get_metrics and get_dense_log (optional)

These methods let your component contribute to environment-level reporting.

def get_metrics(self):
    """
    Returns:
        dict of {"metric_key": scalar_value}, or None to opt out.
    """
    total_bonuses = sum(
        event["bonus"]
        for step in self._bonus_log
        for event in step
    )
    return {"total_bonuses": total_bonuses}

def get_dense_log(self):
    """
    Returns the per-step bonus log accumulated during the episode.
    Return None to opt out of dense logging.
    """
    return self._bonus_log

Metric keys are namespaced by the component’s shorthand when the environment assembles them:

# env.metrics will contain:
# {"Production/total_bonuses": 42, ...}

Registering the component

Decorating the class with @component_registry.add registers it under component.name:

from ai_economist.foundation.base.base_component import (
    BaseComponent,
    component_registry,
)

@component_registry.add
class BonusOutput(BaseComponent):
    name = "BonusOutput"
    ...

The component is then available as:

import ai_economist.foundation as foundation

# Only visible if your module is imported before this call.
BonusOutputClass = foundation.components.get("BonusOutput")

For the registry to find your class, the file that defines it must be imported before foundation.components.get() is called. Add an import to ai_economist/foundation/components/__init__.py or import the module manually at the top of your script.

Complete minimal example

bonus_output.py

import numpy as np
from ai_economist.foundation.base.base_component import (
    BaseComponent,
    component_registry,
)


@component_registry.add
class BonusOutput(BaseComponent):
    """Gives mobile agents a chance to earn bonus Coin each timestep."""

    name = "BonusOutput"
    component_type = "Production"
    agent_subclasses = ["BasicMobileAgent"]
    required_entities = ["Coin", "Labor"]

    def __init__(self, *base_component_args, bonus_rate=0.1, **base_component_kwargs):
        super().__init__(*base_component_args, **base_component_kwargs)
        self.bonus_rate = float(bonus_rate)
        self._bonus_log = []

    def get_n_actions(self, agent_cls_name):
        if agent_cls_name == "BasicMobileAgent":
            return 1
        return None

    def get_additional_state_fields(self, agent_cls_name):
        if agent_cls_name == "BasicMobileAgent":
            return {"bonus_output_count": 0}
        return {}

    def component_step(self):
        step_log = []
        for agent in self.world.agents:
            action = agent.get_component_action(self.name)
            if action == 0:
                continue
            bonus = int(np.random.rand() < self.bonus_rate)
            agent.state["inventory"]["Coin"] += bonus
            agent.state["endogenous"]["Labor"] += 1.0
            agent.state["bonus_output_count"] += bonus
            if bonus:
                step_log.append({"agent": agent.idx, "bonus": bonus})
        self._bonus_log.append(step_log)

    def generate_observations(self):
        return {
            str(agent.idx): {
                "bonus_output_count": agent.state["bonus_output_count"],
                "bonus_rate": self.bonus_rate,
            }
            for agent in self.world.agents
        }

    def generate_masks(self, completions=0):
        masks = {}
        for agent in self.world.agents:
            can_act = 1.0 if agent.state["bonus_output_count"] == 0 else 0.0
            masks[agent.idx] = np.array([can_act], dtype=np.float32)
        return masks

    def additional_reset_steps(self):
        self._bonus_log = []

    def get_metrics(self):
        total = sum(e["bonus"] for step in self._bonus_log for e in step)
        return {"total_bonuses": total}

    def get_dense_log(self):
        return self._bonus_log

Using it in an environment:

import bonus_output  # trigger registration
import ai_economist.foundation as foundation

ScenarioClass = foundation.scenarios.get("uniform/simple_wood_and_stone")

env = ScenarioClass(
    components=[
        ("Gather", {"move_labor": 1.0, "collect_labor": 2.0}),
        ("BonusOutput", {"bonus_rate": 0.15}),
    ],
    n_agents=4,
    world_size=[15, 15],
)

obs = env.reset()

Get Started

Core Concepts

Simulations

Training with RL

Extending Foundation

How components fit into the environment

The Gather component: a worked example

Step-by-step: implementing a custom component

Registering the component

Complete minimal example

Build docs developers (and LLMs) love

Get Started

Core Concepts

Simulations

Training with RL

Extending Foundation

Documentation Index

​How components fit into the environment

​The Gather component: a worked example

​Step-by-step: implementing a custom component

​Registering the component

​Complete minimal example

Build docs developers (and LLMs) love

How components fit into the environment

The Gather component: a worked example

Step-by-step: implementing a custom component

Registering the component

Complete minimal example