Creating custom components

Extend Foundation with new agent-environment dynamics by subclassing BaseComponent.

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:

1. component_step() — execute action effects
2. generate_observations() — collect per-component observations
3. 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

1. 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.

2. 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
   
3. 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
   
4. 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.

5. 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.

6. 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.

7. 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), ...}
   
8. 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 = []
   
9. 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()