tutorials/economic_simulation_advanced.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) 2020, salesforce.com, inc.  \n",
    "All rights reserved.  \n",
    "SPDX-License-Identifier: BSD-3-Clause  \n",
    "For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Colab\n",
    "\n",
    "Try this notebook on [Colab](http://colab.research.google.com/github/salesforce/ai-economist/blob/master/tutorials/economic_simulation_advanced.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Structure of Foundation + How to Extend It \n",
    "\n",
    "In this tutorial, we will explain the low-level compositional structure of Foundation, the economic simulation. Its architecture stems from three main design goals:\n",
    "\n",
    "1. Flexibility: e.g., it should be easy to create worlds with or without income taxes.\n",
    "2. Extensibility: adding new entities and components should follow an easy and transparent process.\n",
    "3. Simplicity: avoid deep class hierarchies.\n",
    "\n",
    "To support these goals, Foundation modularizes the pieces of the simulation as much as possible. Below, we explain what these pieces actually are and how simulation environments are built from them. Understanding this structure will undoubtedly be useful when extending Foundation.\n",
    "\n",
    "Foundation builds economic simulations using Scenario classes. Scenarios compose the simulation's constituent classes into an actual simulation environment. The majority of this tutorial is used to introduce the semantics of each such class type:\n",
    "\n",
    "1. Scenario\n",
    "2. World\n",
    "    - Maps\n",
    "3. Entities\n",
    "    - Resources\n",
    "    - Landmarks\n",
    "    - Endogenous\n",
    "4. Agents\n",
    "5. Components\n",
    "\n",
    "\n",
    "To conclude the tutorial, we will focus on how to extend the economic simulation:\n",
    "\n",
    "6. How the simulation pieces interact\n",
    "7. Exercise: creating a new Component\n",
    "8. Helpful tips"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Before we jump into the details...\n",
    "... let's revisit some basics (covered in detail [here](https://github.com/salesforce/ai-economist/blob/master/tutorials/economic_simulation_basic.ipynb)).\n",
    "\n",
    "Simulation environments exist as Python objects and are interacted with through a gym-style API:\n",
    "```python\n",
    "env = Scenario(...)\n",
    "obs = env.reset()\n",
    "obs, rew, done, info = env.step(actions) # w/ actions <-- policy(obs)\n",
    "```\n",
    "An environment is responsible for providing some *observations* based on the world & agent states and updating these states based on the *actions* taken by the agents and the dynamics of the environment.\n",
    "\n",
    "These dynamics are encapsulated in **Scenario** and **Component** classes, and most extensions of the simulation framework will likely focus on those classes. As a general description...\n",
    "- A **Scenario** provides the backbone of the environment: \n",
    "    - it sets up the world and the agents,\n",
    "    - adds some passive dynamics, \n",
    "    - supplies some observations, \n",
    "    - and generates rewards.\n",
    "- **Components** are how agents interact with the environment: \n",
    "    - they add actions,\n",
    "    - mediate the effect of whatever actions they add,\n",
    "    - and provide relevant observations.\n",
    "\n",
    "(now for the details!)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies\n",
    "\n",
    "You can install the ai-economist package using the pip package manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, signal, sys, time\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "\n",
    "if IN_COLAB:\n",
    "    !pip install ai-economist\n",
    "    \n",
    "    # Restart the Python runtime to automatically use the installed packages\n",
    "    print(\"\\n\\nRestarting the Python runtime! Please (re-)run the cells below.\")\n",
    "    time.sleep(1)\n",
    "    os.kill(os.getpid(), signal.SIGKILL)\n",
    "else:\n",
    "    !pip install ai-economist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai_economist import foundation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Registry\n",
    "\n",
    "The Registry class enables conveniently creating classes, such as Scenarios and Resources, using their names (as a string). Each class can be added to the Registry using the *add()* method and can be retrieved using their *name* property.\n",
    "\n",
    "For example, the *make_env_instance* convenience method used below, can create a basic ```\"layout_from_file/simple_wood_and_stone\"``` Scenario by first retrieving the associated Class from the registry:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai_economist.foundation.base.base_env import BaseEnvironment, scenario_registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_env_cls = scenario_registry.get(\"layout_from_file/simple_wood_and_stone\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_env_cls.name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To add a new class to a Registry, you can use a decorator as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@scenario_registry.add\n",
    "class NewEnvironment(BaseEnvironment):\n",
    "    name = \"NewEnvironment\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_env_cls = scenario_registry.get(\"NewEnvironment\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_env_cls.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are the Scenario classes registered in scenario_registry\n",
    "print(scenario_registry.entries)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is a separate registry for each type of environment component. The ```foundation``` package exposes them as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scenarios:\n",
    "print(foundation.scenarios.entries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Entities (landmarks, resources, endogenous):\n",
    "print(foundation.landmarks.entries)\n",
    "print(foundation.resources.entries)\n",
    "print(foundation.endogenous.entries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Agents:\n",
    "print(foundation.agents.entries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Components:\n",
    "print(foundation.components.entries)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Scenarios\n",
    "\n",
    "As discussed in [the basics tutorial](https://github.com/salesforce/ai-economist/blob/master/tutorials/economic_simulation_basic.ipynb), the Scenario class implements an economic simulation with multiple agents and (optionally) a social planner. \n",
    "\n",
    "We will create the same environment instance used in that tutorial, using the configuration below. This configuration defines a simulation with four agents in a world of 15 by 15 cells. Each Agent can: \n",
    "\n",
    "- gather collectible Resources (through the Gather Component), \n",
    "- build Houses (through the Build Component), and \n",
    "- trade collectible Resources (through the ContinuousDoubleAuction Component)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the configuration of the environment that will be built\n",
    "\n",
    "env_config = {\n",
    "    # ===== STANDARD ARGUMENTS ======\n",
    "    'n_agents': 4,          # Number of non-planner agents\n",
    "    'world_size': [15, 15], # [Height, Width] of the env world\n",
    "    'episode_length': 1000, # Number of timesteps per episode\n",
    "    \n",
    "    # In multi-action-mode, the policy selects an action for each action subspace (defined in component code)\n",
    "    # Otherwise, the policy selects only 1 action\n",
    "    'multi_action_mode_agents': False,\n",
    "    'multi_action_mode_planner': True,\n",
    "    \n",
    "    # When flattening observations, concatenate scalar & vector observations before output\n",
    "    # Otherwise, return observations with minimal processing\n",
    "    'flatten_observations': False,\n",
    "    # When Flattening masks, concatenate each action subspace mask into a single array\n",
    "    # Note: flatten_masks = True is recommended for masking action logits\n",
    "    'flatten_masks': True,\n",
    "    \n",
    "    \n",
    "    # ===== COMPONENTS =====\n",
    "    # Which components to use (specified as list of {\"component_name\": {component_kwargs}} dictionaries)\n",
    "    #   \"component_name\" refers to the component class's name in the Component Registry\n",
    "    #   {component_kwargs} is a dictionary of kwargs passed to the component class\n",
    "    # The order in which components reset, step, and generate obs follows their listed order below\n",
    "    'components': [\n",
    "        # (1) Building houses\n",
    "        {'Build': {}},\n",
    "        # (2) Trading collectible resources\n",
    "        {'ContinuousDoubleAuction': {'max_num_orders': 5}},\n",
    "        # (3) Movement and resource collection\n",
    "        {'Gather': {}},\n",
    "    ],\n",
    "    \n",
    "    # ===== SCENARIO =====\n",
    "    # Which scenario class to use (specified by the class's name in the Scenario Registry)\n",
    "    'scenario_name': 'uniform/simple_wood_and_stone',\n",
    "    \n",
    "    # (optional) kwargs of the chosen scenario class\n",
    "    'starting_agent_coin': 10,\n",
    "    'starting_stone_coverage': 0.10,\n",
    "    'starting_wood_coverage':  0.10,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This configuration dictionary lists the used Components, each can be configured through a dictionary of Component-specific settings. \n",
    "\n",
    "Creating a Scenario can be done using a convenience method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = foundation.make_env_instance(**env_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "obs = env.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above code, ```env``` is an instance of the Scenario class stored in ```scenario_registry``` as ```\"uniform/simple_wood_and_stone\"```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uniform_cls = scenario_registry.get(env_config['scenario_name'])\n",
    "isinstance(env, uniform_cls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This Scenario class (and all Scenario classes) are subclasses of ```BaseEnvironment``` (meaning ```env``` is also an instance of ```BaseEnvironment```)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "isinstance(env, BaseEnvironment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Why this structure?** The ```env``` object is responsible for a lot! It organizes all the pieces (the world, agents, and components) into a coherent environment with a simple and consistent API. It also implements some of the behavior of the environment itself: the passive (not action-dependent) dynamics of the world, baseline observations, and rewards.\n",
    "\n",
    "That first domain of functionality is implemented in the ```BaseEnvironment``` code and the second domain of functionality (the \"behavior\") is implemented separately by each Scenario class via the following methods:\n",
    "```python\n",
    "from ai_economist.foundation.base.base_env import BaseEnvironment, scenario_registry\n",
    "\n",
    "@scenario_registry.add\n",
    "class EmptyScenario(BaseEnvironment):\n",
    "    name = \"Empty\"\n",
    "    required_entities = []\n",
    "    \n",
    "    def reset_layout(self):\n",
    "        \"\"\"Resets the state of the world object (self.world).\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def reset_agent_states(self):\n",
    "        \"\"\"Resets the state of the agent objects (self.world.agents & self.world.planner).\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def scenario_step(self):\n",
    "        \"\"\"Implements the passive dynamics of the environment.\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def generate_observations(self):\n",
    "        \"\"\"Yields some basic observations about the world/agent states.\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def compute_reward(self):\n",
    "        \"\"\"Determines the reward each agent receives at the end of each timestep.\"\"\"\n",
    "        pass\n",
    "```\n",
    "\n",
    "The expected behaviors of these methods are described extensively in the internal documentation of ```BaseEnvironment```, where they are defined as abstract methods (see [foundation/base/base_env.py](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/base/base_env.py))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```env```, which is an instance of the ```\"uniform/simple_wood_and_stone\"``` Scenario class, has the following behavior:\n",
    "- **reset_layout**: Samples a new spatial layout of Stone and Wood source tiles in the world.\n",
    "- **reset_agent_states**: Resets agent inventories and their starting locations in the world.\n",
    "- **scenario_step**: Stochastically re-spawns Stone and Wood at empty source tiles.\n",
    "- **generate_observations**: Generates observations related to inventory and the spatial state of the world.\n",
    "- **compute_reward**: Marginal utility for each agent in ```env.world.agents```; marginal social welfare for ```env.world.planner```.\n",
    "\n",
    "Check out [the code for this Scenario class](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/scenarios/simple_wood_and_stone/dynamic_layout.py) to see how this behavior is implemented.\n",
    "\n",
    "**Note**: This example refers to some concepts we haven't introduced yet (source tiles, inventories, etc.). We'll cover those in the sections below!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. World and Maps\n",
    "\n",
    "Above, we saw how a Scenario class resets the spatial state of the world, but **where is this spatial state represented?**\n",
    "\n",
    "Each Scenario will include an instance of the **World** class ```env.world``` to wrap agent instances (more on that below) and an instance of the **Maps** class ```env.world.maps```. **The Maps class stores and manipulates the spatial state of the environment**, such as the locations of Agents and other Entities.\n",
    "\n",
    "Both classes (World and Maps) are implemented in [foundation/base/world.py](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/base/world.py).\n",
    "\n",
    "The **maps** object ```env.world.maps``` holds a 2-dimensional NumPy array that records the location of Entities in the world."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For each key, the maps object has a [Height, Width] array for the spatial layout of that Entity in the world.\n",
    "env.world.maps.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For instance, we can visualize where the Stone is in the world:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: this map has the same size as our world (15 by 15)\n",
    "env.world.maps.get(\"Stone\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **world** object ```env.world``` provides some tools for interfacing with **maps**.\n",
    "\n",
    "To see which Resources are in a certain cell, you can use the convenience method *location_resources*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.world.location_resources(0, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see which Landmarks are in a certain cell, you can use the convenience method *location_landmarks*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.world.location_landmarks(0, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Entities\n",
    "\n",
    "Agents in the economic simulation can interact with Entities. There are 3 groups of Entities, each with their own semantics:\n",
    "\n",
    "- **Landmarks** show up in the spatial world \n",
    "- **Resources** show up in agent inventories and (optionally) also the spatial world\n",
    "- **Endogenous** entities represent abstract quantities (like effort) that agents can only observe about themselves\n",
    "\n",
    "You can find their definitions in [foundation/entities](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/entities). Again, we will use convenient Registries to retrieve the various classes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Landmarks\n",
    "\n",
    "Landmarks represent entities that exist exclusively in the spatial world, for example a block of Water that agents can't move over.\n",
    "\n",
    "In the current implementation, there are three types of Landmarks: House, Water and SourceBlock. SourceBlocks are special Landmarks from which Resources can spawn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "house = foundation.landmarks.get(\"House\")\n",
    "water = foundation.landmarks.get(\"Water\")\n",
    "source_block_wood = foundation.landmarks.get(\"WoodSourceBlock\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each Landmark, the class defines:\n",
    "\n",
    "- its name, \n",
    "- its color, \n",
    "- whether the Landmark is ownable by an Agent (e.g., Houses), and \n",
    "- whether it's solid (e.g., Water).\n",
    "\n",
    "An agent cannot occupy the same location as a solid landmark (e.g. Water) unless it owns that landmark (e.g. a House)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[k for k in dir(house) if k[0] != \"_\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: The simulation does *not* instantiate a separate Python instance of a Landmark for each occurrence of a Landmark in the world! Rather, each Landmark class defines the abstract properties of any instance of that Landmark.\n",
    "\n",
    "```env.world.maps``` keeps track of where all the units of a Landmark are using a 2-dimensional NumPy array, as illustrated above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resources\n",
    "\n",
    "Resources are another important type of Entity in the world. Resources are semantically different from Landmarks, because Resources can be traded, collected, and converted into other Entities (e.g., Wood and Stone are used to build a House).\n",
    "\n",
    "In particular, Resources are the entities in the world that an agent can own as part of its **inventory**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.get_agent(agent_idx=0).inventory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A Resource has three main attributes: \n",
    "- its name, \n",
    "- its color (convenient for visualization),\n",
    "- and whether it's collectible.\n",
    "\n",
    "For example, we can see that Wood is collectible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wood = foundation.resources.get(\"Wood\")\n",
    "\n",
    "wood.name, wood.collectible, wood.color"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the other hand, Coin is *not* collectible (but can be owned)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coin = foundation.resources.get(\"Coin\")\n",
    "\n",
    "coin.collectible"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that collectible Resources (Wood & Stone) get a special Landmark type (source blocks), which both show up in the Map, whereas non-collectible Resources (Coin) only exist as part of the inventory.\n",
    "\n",
    "In other words, **collectible Resources start as part of the spatial world but can be moved into an agent's inventory.**\n",
    "\n",
    "```env.world.maps``` keeps track of where all the units of a **collectible** Resource are using a 2-dimensional NumPy array, as illustrated above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Endogenous Entities\n",
    "\n",
    "Certain semantic concepts do not have a physical realization, but are important because they determine, e.g., an Agent's utility. The main example is Labor.\n",
    "\n",
    "The definition of Labor is rather simple, it only defines the name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labor = foundation.endogenous.get(\"Labor\")\n",
    "\n",
    "[k for k in dir(labor) if k[0] != \"_\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Endogenous entities, like Resources, can be accumulated, **but their quantities are stored outside of the inventory**. (This is done to make it easier to separate Resources and Endogenous entities when generating observations.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent0 = env.get_agent(agent_idx=0)\n",
    "print(agent0.inventory)\n",
    "print(agent0.endogenous)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Agents\n",
    "\n",
    "```env.world``` also wraps **agent** instances. There will be ```env.n_agents``` \"mobile\" agents + 1 \"planner\" agent. Each such agent is represented as a separate Python object:\n",
    "\n",
    "The ```env.n_agents``` \"mobile\" agents (representing individual workers in the economy):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.world.agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The \"planner\" agent (representing a Social Planner that sets, for example, tax policy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.world.planner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Agents can be easily accessed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent0  = env.get_agent(agent_idx=0)   # Mobile agents are numerically indexed\n",
    "agent1  = env.get_agent(agent_idx=1)\n",
    "planner = env.get_agent(agent_idx='p') # The planner agent always uses index 'p'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each agent instance maintains the state of the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent0.state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent1.state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "planner.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Components\n",
    "\n",
    "Up to this point, we have learned about how the state of the world is represented in the **world** object: with spatial state represented by ```env.world.maps```, and agent states represented in ```env.world.agents``` and ```env.world.planner```.\n",
    "\n",
    "We have also learned how custom **Scenario** classes define methods for resetting these states and rules for passive dynamics (in our working example, resource regeneration).\n",
    "\n",
    "**How then do agents actually _interact_ with the environment?**\n",
    "\n",
    "**Components** are used to flexibly extend the behavior of a Scenario by encapsulating specific interactions/dynamics. They enable a plug-and-play approach to building economic simulations.\n",
    "\n",
    "This structure also vastly simplifies the process of extending the simulation framework through the addition of new Component classes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's revisit our working example to better understand how Components work...\n",
    "\n",
    "... recall the ```'components'``` argument set in the environment configuration we used to build ```env```:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_config['components']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This argument tells the Scenario class which Component classes to make use of. Notice that ```env``` has created an instance of each such class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env._components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# which are better accessed via...\n",
    "build = env.get_component(\"Build\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```build``` is an instance of the Component class ```Build```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "isinstance(build, foundation.components.get(\"Build\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All Component classes are subclasses of ```BaseComponent``` (so ```build``` is also an instance of ```BaseComponent```)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai_economist.foundation.base.base_component import BaseComponent\n",
    "isinstance(build, BaseComponent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Why this structure?** Building Component classes (such as ```Build```) on top of ```BaseComponent``` enforces consistent semantics for defining the behavior of the Component and allowing a Scenario to make use of it. Each Component class must implement the following abstract methods:\n",
    "\n",
    "```python\n",
    "from foundation.base.base_component import BaseComponent, component_registry\n",
    "\n",
    "@component_registry.add\n",
    "class EmptyComponent(BaseComponent):\n",
    "    name = \"Empty\"\n",
    "    required_entities = []\n",
    "    \n",
    "    def get_n_agent_actions(self, agent_cls_name):\n",
    "        \"\"\"Returns the actions that agents with type agent_cls_name can take through this component.\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def get_additional_state_fields(self, agent_cls_name):\n",
    "        \"\"\"Returns a dictionary to be be added to the state dictionary of agents with type agent_cls_name.\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def component_step(self):\n",
    "        \"\"\"Implements the (passive and active) dynamics that this Component adds to the environment.\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def generate_observations(self):\n",
    "        \"\"\"Yields observations.\"\"\"\n",
    "        pass\n",
    "    \n",
    "    def generate_masks(self):\n",
    "        \"\"\"Specifies which of the Component actions are valid given the current state.\"\"\"\n",
    "        pass\n",
    "```\n",
    "\n",
    "The expected behaviors of these methods are described extensively in the internal documentation of ```BaseComponent```, where they are defined as abstract methods (see [foundation/base/base_component.py](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/base/base_component.py))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an example, the ```Build``` Component class implements the following behavior:\n",
    "- **get_n_agent_actions**: returns 1 action that mobile agents can take (to build a house).\n",
    "- **get_additional_state_fields**: returns mobile agents' state dictionary, which includes payment-per-house info.\n",
    "- **component_step**: For each agent that takes the build action, place an agent-owned house landmark at the agent's location and update its state (remove Stone & Wood used for building, add Coin income, add Labor cost).\n",
    "- **generate_observations**: Generates observations related to the payment-per-house state info.\n",
    "- **generate_masks**: Mask the build action for agents that are on non-empty map cells or do not have the resources to build.\n",
    "\n",
    "Check out [the code for the Build Component class](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/components/build.py) to see how this behavior is implemented."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an additional example, the ```Gather``` Component class implements the following behavior:\n",
    "- **get_n_agent_actions**: returns 4 actions that mobile agents can take (move up, down, left, and right).\n",
    "- **get_additional_state_fields**: returns mobile agents' state dictionary, which includes probability of collecting bonus resources.\n",
    "- **component_step**: For each agent that takes a move action: update its location, if the new location has a Resource, move it from the spatial world to the agent's inventory, add Labor cost(s) associated with moving and collecting.\n",
    "- **generate_observations**: Generates observations related to the bonus probability state info.\n",
    "- **generate_masks**: For each agent, mask whichever move actions would move it to a location it is not allowed to occupy.\n",
    "\n",
    "Check out [the code for the Gather Component class](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/components/move.py) to see how this behavior is implemented."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the agent objects are created, each one registers the actions afforded to it by the ```env``` components. As described: ```Build``` adds 1 action and ```Gather``` adds 4. ```ContinuousDoubleAuction``` (which implements trading) is more complex: it adds several action sets for buying and selling Wood and Stone (each action in a set corresponds to a different price level).\n",
    "\n",
    "More concretely, the ```ContinuousDoubleAuction``` Component class implements the following behavior:\n",
    "- **get_n_agent_actions**: (For mobile agents) returns a *pair* action action sets (for buying and selling) for each *collectible* resource; each action set adds M+1 actions, where M is the maximum trading price.\n",
    "- **get_additional_state_fields**: Doesn't add any state fields.\n",
    "- **component_step**: For each agent that takes a buy/sell action: create an order in the associated resource market and add a small Labor cost; match orders and execute trades, which moves Coin and the resource between inventories; update the order books, removing expired orders.\n",
    "- **generate_observations**: For each agent, generates observations related to price levels of past successful trades, current available orders, and the agent's own outstanding orders.\n",
    "- **generate_masks**: For each agent, mask any buying actions that it does not have enough Coin to fulfill and mask any selling actions that it does not have the resources to fulfill.\n",
    "\n",
    "Check out [the code for the ContinuousDoubleAuction Component class](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/components/continuous_double_auction.py) to see how this behavior is implemented."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After setting up the environment, the agents have registered the following actions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent0.action_dim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Right now, the planner has not registered any actions because none of the 3 Components add any planner actions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "planner.action_dim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In particular, we did not include a Taxation Component in ```env_config['components']```, which would introduce actions for the planner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "tax_config = deepcopy(env_config)\n",
    "tax_config['components'].append({\"PeriodicBracketTax\": {}})\n",
    "\n",
    "tax_env = foundation.make_env_instance(**tax_config)\n",
    "\n",
    "tax_env._components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same as with the other env (PeriodicBracketTax doesn't add actions for this agent type)\n",
    "tax_env.get_agent(agent_idx=0).action_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now the planner has actions (PeriodicBracketTax creates an action set for the planner for each tax bracket)\n",
    "tax_env.get_agent(agent_idx='p').action_dim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extending the Economic Simulation\n",
    "\n",
    "Having introduced the constituent parts of the simulation, we now focus on how to extend it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. How the Simulation Pieces Interact\n",
    "\n",
    "Before we can dive into actually writing new code, we still need to understand how the simulation pieces interact to build a coherent environment with a consistent and simple API.\n",
    "\n",
    "In particular, let's look at how the environment sets itself up when an environment instance is created and what happens under the hood when stepping through the simulation.\n",
    "\n",
    "The underlying design follows the principle of *encapsulation* by making it easy to create new simulation classes, such as Entities, Scenarios, and Components, without having to re-write existing code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up\n",
    "\n",
    "Earlier, when we looked at ```env.world.maps```, we saw that it already had populated a handful of maps, ones for Wood, Stone, their associated SourceBlocks, and Houses.\n",
    "\n",
    "Also, when we looked at some of the ```env.world.agents```, we saw that their states were already populated, for example with inventories referencing Resources such as Wood, Stone, and Coin.\n",
    "\n",
    "**Where did that come from?**\n",
    "\n",
    "Each Scenario and Component must declare the entities that it will interact with in the attribute ```required_entities```. For example:\n",
    "\n",
    "```python\n",
    "@scenario_registry.add\n",
    "class Uniform(BaseEnvironment):\n",
    "    name = \"uniform/simple_stone_and_wood\"\n",
    "    required_entities = [\"Stone\", \"Wood\"]\n",
    "    ...\n",
    "```\n",
    "\n",
    "or\n",
    "\n",
    "```python\n",
    "@component_registry.add\n",
    "class Build(BaseComponent):\n",
    "    name = \"Build\"\n",
    "    required_entities = [\"Stone\", \"Wood\", \"House\"]\n",
    "    ...\n",
    "```\n",
    "\n",
    "When ```env``` gets created (as part of ```BaseEnvironment.__init__```), the following happens:\n",
    "\n",
    "1. It looks at the ```required_entities``` of the Scenario class and the included Component classes and it determines which Resources, Landmarks, and Endogenous entities need to be part of the game.\n",
    "    - By default, Coin and Labor are always included, even if they are not mentioned in the Scenario's or Components' ```required_entities```.\n",
    "\n",
    "\n",
    "2. It constructs a world object ```env.world```, which involves creating partially-initialized agent objects for each agent in the environment and creating the maps object ```env.world.maps```. \n",
    "    - The world object is told which Resources and Landmarks to include and this gets passed to the maps object. The maps object creates a map for these and uses their class properties to preserve semantics: for example, the maps object maintain a separate map indicating ownership for ownable Landmarks, like Houses.\n",
    "    - Agent objects are also told which Resources and Endogenous entities are in use. The ```inventory```, ```escrow```, and ```endogenous``` portions of their state dictionaries are configured accordingly.\n",
    "\n",
    "\n",
    "3. It creates an instance of each of the included Component classes (using, for each, any paired keyword arguments).\n",
    "\n",
    "\n",
    "4. It finishes initializing the agent objects by allowing each one to register the actions defined by the different component objects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stepping\n",
    "\n",
    "Outside of initialization, the logic for integrating Scenarios with Components is fairly straightforward.\n",
    "\n",
    "When calling ```env.step(actions)``` (the main method for interacting with the environment), the following happens:\n",
    "1. The environment interprets ```actions```, which updates each agent objects' ```agent.action``` to represent which action the agent is taking for each of the action sets (0 denotes NO-OP, meaning no action).\n",
    "\n",
    "\n",
    "2. The environment performs the ```component_step``` method for each of the component objects. Inside ```component_step```, agent method ```agent.get_component_action(self.name, [action_set_name])``` can be used to query the action(s) ```agent``` chose.\n",
    "\n",
    "\n",
    "3. The environment performs the ```scenario_step``` of its own Scenario class.\n",
    "\n",
    "\n",
    "4. For each agent, ```agent.action``` is reset to all NO-OPs.\n",
    "\n",
    "\n",
    "5. The environment collects observations from its own ```generate_observations``` method and those of each of the component objects, and it combines and formats them into a single observations dictionary.\n",
    "\n",
    "\n",
    "6. The environment collects action masks using the ```generate_masks``` method of each of the component objects, and it combines and formats them before packaging them as ```'action_masks'``` in each agent's observations.  \n",
    "  \n",
    "  \n",
    "\n",
    "\n",
    "A similar logic is applied for ```env.reset()```, in which:\n",
    "1. ```env.reset_starting_layout``` and ```env.reset_agent_states``` (which are defined by the Scenario class) are first called.\n",
    "\n",
    "\n",
    "2. Each component object's ```reset``` method is called.\n",
    "\n",
    "\n",
    "3. Finally, ```env.additional_reset_steps``` is called (which is also defined by the Scenario class)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7. Exercise: Creating a New BuyWidgetFromVirtualStore Component\n",
    "\n",
    "Let's put all these concepts together and introduce a new Resource entity, a Widget, and implement a simple Component in which agents can buy Widgets from an external source, like an online store, for a fixed price of 5 Coin. The store randomly adds a single Widget to its inventory each step. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding \"Widget\" as a new Resource entity\n",
    "\n",
    "In order to add a new Resource entity that other Scenario and Component classes can reference, we simply need to define a new Resource class and put it in the appropriate registry. Let's do this directly in code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai_economist.foundation.entities.resources import Resource, resource_registry\n",
    "\n",
    "@resource_registry.add\n",
    "class Widget(Resource):\n",
    "    name = \"Widget\"\n",
    "    color = [1, 1, 1]\n",
    "    collectible = False # <--- Goes in agent inventory, but not in the world"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it. It's that easy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Component Initialization\n",
    "\n",
    "Let's start with the initialization of the Component. We'll define a customizable ```widget_refresh_rate``` which determines how likely the store will add a new Widget unit to its inventory each step. Additionally, we'll use a fixed price of 5 Coin per Widget, and initialize the store's inventory to 0.\n",
    "\n",
    "```python\n",
    "@component_registry.add\n",
    "class BuyWidgetFromVirtualStore(BaseComponent):\n",
    "    name = \"BuyWidgetFromVirtualStore\"\n",
    "    required_entities = [\"Coin\", \"Widget\"]  # <--- We can now look up \"Widget\" in the resource registry\n",
    "    agent_subclasses = [\"BasicMobileAgent\"]\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        *base_component_args,\n",
    "        widget_refresh_rate=0.1,\n",
    "        **base_component_kwargs\n",
    "    ):\n",
    "        super().__init__(*base_component_args, **base_component_kwargs)\n",
    "        self.widget_refresh_rate = widget_refresh_rate\n",
    "        self.available_widget_units = 0\n",
    "        self.widget_price = 5\n",
    "```        \n",
    "\n",
    "Note that we define the Component's name as a string ```BuyWidgetFromVirtualStore```, and decorate the Component with the ```add``` method from the component registry. This allows us to create the Component by using the ```get``` method on ```component_registry```.\n",
    "\n",
    "We also declare the ```required_entities``` as ```Coin``` and ```Widget```. This instructs ```BaseEnvironment``` to include these entity types in the environment when ```BuyWidgetFromVirtualStore``` is used as a component."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reset\n",
    "\n",
    "Sometimes, a Component wants to expose part of the state it manages as a part of the agents' state. ```get_additional_state_fields``` is used to set that up and reset the associated state when a new episode starts. Here, we won't use that functionality so we return an empty dictionary, which is interpreted as *no additional state fields*.\n",
    "\n",
    "```python\n",
    "def get_additional_state_fields(self, agent_cls_name):\n",
    "    return {}\n",
    "```\n",
    "\n",
    "When a new episode starts (whenever the ```BaseEnvironment``` resets), the store should have 0 Widgets. We can use ```additional_reset_steps``` to implement this behavior.\n",
    "\n",
    "```python    \n",
    "def additional_reset_steps(self):\n",
    "    self.available_wood_units = 0\n",
    "```        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Actions\n",
    "\n",
    "Each agent can choose to *buy* a Widget or not each step. Hence, we add an extra action to the action space of a ```BasicMobileAgent```. Other agent types (like planners) do not get an extra action: if ```get_n_actions``` returns ```None```, it is interpreted as *no action added*.  **Note**: the simulation framework only supports discrete action types for now.\n",
    "\n",
    "```python\n",
    "def get_n_actions(self, agent_cls_name):\n",
    "    # This component adds 1 binary action that mobile agents can take: buy widget (or not).\n",
    "    if agent_cls_name == \"BasicMobileAgent\":\n",
    "        return 1  # Buy or not.\n",
    "\n",
    "    return None\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Action Masks\n",
    "\n",
    "Whether or not an agent can buy depends on:\n",
    "\n",
    "- Does the agent have at least ```self.widget_price``` Coin? We check this by looking at ```agent.state[\"inventory\"][\"Coin\"]```.\n",
    "- Does the store have at least 1 Widget in store?\n",
    "\n",
    "Because a BasicMobileAgent has 1 extra discrete action, the mask is simply a single bit, stored in a NumPy array.\n",
    "\n",
    "**Note: ```world.agents``` loops over ```BasicMobileAgent```s only! It does not include the planner agent!**\n",
    "\n",
    "```python\n",
    "def generate_masks(self, completions=0):\n",
    "    masks = {}\n",
    "    # Mobile agents' buy action is masked if they cannot build with their\n",
    "    # current coin or if no widgets are available.\n",
    "    for agent in self.world.agents:\n",
    "        masks[agent.idx] = np.array([\n",
    "            agent.state[\"inventory\"][\"Coin\"] >= self.widget_price and self.available_widget_units > 0\n",
    "        ])\n",
    "\n",
    "    return masks\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step\n",
    "\n",
    "The main logic of this Component is defined in ```component_step```. Two pieces of logic are defined: \n",
    "\n",
    "1. The store randomly adds a unit of Wood to its inventory. \n",
    "2. Agents buy orders are executed in random order (to break ties if, say, there's only 1 Widget but 2 agents try to buy it).\n",
    "\n",
    "```python\n",
    "def component_step(self):\n",
    "    # Maybe add a Widget to store's inventory.\n",
    "    if random.random() < self.widget_refresh_rate: \n",
    "        self.available_widget_units += 1\n",
    "\n",
    "    # Agents can buy 1 unit of Wood, in random order.\n",
    "    for agent in self.world.get_random_order_agents():\n",
    "\n",
    "        action = agent.get_component_action(self.name)\n",
    "        \n",
    "        if action == 0: # NO-OP. Agent is not interacting with this component.\n",
    "            continue\n",
    "\n",
    "        if action == 1: # Agent wants to buy. Execute a purchase if possible.\n",
    "            if self.available_widget_units > 0 and agent.state[\"inventory\"][\"Coin\"] >= self.widget_price: \n",
    "                # Remove the purchase price from the agent's inventory\n",
    "                agent.state[\"inventory\"][\"Coin\"] -= self.widget_price\n",
    "                # Add a Widget to the agent's inventory\n",
    "                agent.state[\"inventory\"][\"Widget\"] += 1\n",
    "                # Remove the Widget from the market\n",
    "                self.available_widget_units -= 1\n",
    "                \n",
    "        else: # We only declared 1 action for this agent type, so action > 1 is an error.\n",
    "            raise ValueError\n",
    "```\n",
    "\n",
    "**Note how the step logic supports action=0 and action=1.** action=0 denotes ``NO-OP`` (no operation). **All Components are expected to obey this semantic.** The action added by this Component starts (and in this case ends) with action=1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Observations\n",
    "\n",
    "The store is quite transparent: each ```BasicMobileAgent``` can observe how likely it is that the store will add new a new Widget unit, what the store's current inventory looks like, and what the price is.\n",
    "\n",
    "The observation that a Component generates should be structured as a dictionary, keyed by each agent's ```id``` and each value being a dictionary.\n",
    "\n",
    "```python\n",
    "def generate_observations(self):\n",
    "    obs_dict = dict()\n",
    "    for agent in self.world.agents:\n",
    "        obs_dict[agent.idx] = {\n",
    "            \"widget_refresh_rate\": self.widget_refresh_rate,\n",
    "            \"available_widget_units\": self.available_widget_units,\n",
    "            \"widget_price\": self.widget_price\n",
    "        }\n",
    "\n",
    "    return obs_dict\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Final Component\n",
    "\n",
    "Let's combine this into actual code so we can create the new Component class and have it available in the component registry:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from ai_economist.foundation.base.base_component import BaseComponent, component_registry\n",
    "\n",
    "@component_registry.add\n",
    "class BuyWidgetFromVirtualStore(BaseComponent):\n",
    "    name = \"BuyWidgetFromVirtualStore\"\n",
    "    required_entities = [\"Coin\", \"Widget\"]  # <--- We can now look up \"Widget\" in the resource registry\n",
    "    agent_subclasses = [\"BasicMobileAgent\"]\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        *base_component_args,\n",
    "        widget_refresh_rate=0.1,\n",
    "        **base_component_kwargs\n",
    "    ):\n",
    "        super().__init__(*base_component_args, **base_component_kwargs)\n",
    "        self.widget_refresh_rate = widget_refresh_rate\n",
    "        self.available_widget_units = 0\n",
    "        self.widget_price = 5\n",
    "\n",
    "    def get_additional_state_fields(self, agent_cls_name):\n",
    "        return {}\n",
    "\n",
    "    def additional_reset_steps(self):\n",
    "        self.available_wood_units = 0\n",
    "\n",
    "    def get_n_actions(self, agent_cls_name):\n",
    "        if agent_cls_name == \"BasicMobileAgent\":\n",
    "            return 1\n",
    "        return None\n",
    "\n",
    "    def generate_masks(self, completions=0):\n",
    "        masks = {}\n",
    "        for agent in self.world.agents:\n",
    "            masks[agent.idx] = np.array([\n",
    "                agent.state[\"inventory\"][\"Coin\"] >= self.widget_price and self.available_widget_units > 0\n",
    "            ])\n",
    "\n",
    "        return masks\n",
    "\n",
    "    def component_step(self):\n",
    "        if random.random() < self.widget_refresh_rate: \n",
    "            self.available_widget_units += 1\n",
    "\n",
    "        for agent in self.world.get_random_order_agents():\n",
    "\n",
    "            action = agent.get_component_action(self.name)\n",
    "\n",
    "            if action == 0: # NO-OP. Agent is not interacting with this component.\n",
    "                continue\n",
    "\n",
    "            if action == 1: # Agent wants to buy. Execute a purchase if possible.\n",
    "                if self.available_widget_units > 0 and agent.state[\"inventory\"][\"Coin\"] >= self.widget_price: \n",
    "                    agent.state[\"inventory\"][\"Coin\"] -= self.widget_price\n",
    "                    agent.state[\"inventory\"][\"Widget\"] += 1\n",
    "                    self.available_widget_units -= 1\n",
    "\n",
    "            else: # We only declared 1 action for this agent type, so action > 1 is an error.\n",
    "                raise ValueError\n",
    "\n",
    "    def generate_observations(self):\n",
    "        obs_dict = dict()\n",
    "        for agent in self.world.agents:\n",
    "            obs_dict[agent.idx] = {\n",
    "                \"widget_refresh_rate\": self.widget_refresh_rate,\n",
    "                \"available_widget_units\": self.available_widget_units,\n",
    "                \"widget_price\": self.widget_price\n",
    "            }\n",
    "\n",
    "        return obs_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a new environment instance that uses the new Component \n",
    "\n",
    "To add the ```BuyWoodFromVirtualStore``` to the Scenario, modify the ```env_config``` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_env_config = deepcopy(env_config)\n",
    "\n",
    "# Compared to env_config, new_env_config simply adds our new Component\n",
    "new_env_config['components'] = [\n",
    "    # (1) Building houses\n",
    "    {'Build': {}},\n",
    "    # (2) Trading collectible resources\n",
    "    {'ContinuousDoubleAuction': {'max_num_orders': 5}},\n",
    "    # (3) Movement and resource collection\n",
    "    {'Gather': {}},\n",
    "    # (4) Let each mobile agent buy widgets from a virtual store.\n",
    "    {'BuyWidgetFromVirtualStore': {'widget_refresh_rate': 0.1}},  # <--- This.\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And there you have it!\n",
    "new_env = foundation.make_env_instance(**new_env_config)\n",
    "obs = new_env.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare ```env``` and ```new_env```!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_env.resources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how ```new_env``` now includes ```'Widget'``` as a Resource in the environment. This is because ```BuyWidgetFromVirtualStore.required_entities``` includes ```'Widget'```!\n",
    "\n",
    "This difference also shows up in the agent states -- specifically, the inventory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_agent0 = env.get_agent(agent_idx=0)\n",
    "new_agent0 = new_env.get_agent(agent_idx=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inventory includes Coin, Stone, and Wood...\n",
    "old_agent0.state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inventory includes Coin, Stone, Wood and Widget!\n",
    "new_agent0.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mobile agents in ```new_env``` should also have an extra action set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_agent0.action_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_agent0.action_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_agent0.get_component_action('BuyWidgetFromVirtualStore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, with that, ```BuyWidgetFromVirtualStore``` is a brand new Component, ready to go. Pretty cool, huh?!\n",
    "\n",
    "If you're interested in learning more about how to extend the simulation by adding new classes, we encourage you to check out the existing implementations provided in the code and to refer back to the documentation in the base classes on which everything is built!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One last thing (because it's cool)...\n",
    "Once we included our new Component, Widgets automatically became part of the agents' inventory space. That's because we defined the Widget entity as a Resource class.\n",
    "\n",
    "However, Widgets are not part of the spatial map. That's because we defined ```Widget.collectible = False```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# No Widget map:\n",
    "new_env.world.maps.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's re-define the Widget Resource class, but with ```Widget.collectible = True```, which will give it the same semantics as Wood and Stone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai_economist.foundation.entities.landmarks import Landmark, landmark_registry\n",
    "\n",
    "@resource_registry.add\n",
    "class Widget(Resource):\n",
    "    name = \"Widget\"\n",
    "    color = [1, 1, 1]\n",
    "    collectible = True # <--- Goes in agent inventory, AND in the world\n",
    "    \n",
    "# Since we're doing this in a notebook, we need to manually add a Source Block Landmark for Widgets.\n",
    "#     If we defined the Widget class in /foundation/entities/resources.py,\n",
    "#     this class construction would happen automatically.\n",
    "@landmark_registry.add\n",
    "class SourceBlock(Landmark):\n",
    "    \"\"\"Special Landmark for generating collectible resources. Not ownable. Not solid.\"\"\"\n",
    "\n",
    "    name = \"{}SourceBlock\".format(Widget.name)\n",
    "    color = np.array(Widget.color)\n",
    "    ownable = False\n",
    "    solid = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've given Widget new semantics, let's make another environment object and look at the map keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_env_with_collectible_widgets = foundation.make_env_instance(**new_env_config)\n",
    "\n",
    "new_env_with_collectible_widgets.world.maps.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cool! Spatial maps for ```'Widget'``` and ```'WidgetSourceBlock'``` are automatically created because we've defined Widget as something that should be collectible from the spatial world. \n",
    "\n",
    "The Component ```BuyWidgetFromVirtualStore``` will still work just the same -- no need to re-write that. However, the two new maps will always be empty because our Scenario class only handles populating/regenerating Wood and Stone.\n",
    "\n",
    "That's fine. After all, this was just to demonstrate the plug-and-play design of the simulation framework!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 8. Helpful Tips"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Components can be passive.\n",
    "\n",
    "If you wish to introduce a dynamic to the environment that doesn't depend on agent actions, you can do so through a Component class. Components don't need to add actions and the ```component_step``` doesn't need to depend on actions. An example is found in the [WealthRedistribution](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/components/redistribution.py) class."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Components can be stateful.\n",
    "\n",
    "In the actual environment, components are Python objects, so you might as well use them that way. A simple demonstration of this is actually found in the example Component above (```BuyWidgetFromVirtualStore```). Notice how the class has a ```available_widget_units``` attribute, which it creates during ```__init___```, updates during ```component_step```, and resets in ```additional_reset_steps```.\n",
    "\n",
    "Conceptually, ```BuyWidgetFromVirtualStore.available_widget_units``` is just as much a part of the environment state as, say, ```new_agent0.state```. You should feel free to take advantage of the object-oriented design. Just make sure to properly reset internally managed states in ```additional_reset_steps```!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Components can add multiple action sets per agent.\n",
    "\n",
    "Components can create many sets of actions. For example (for mobile agents):\n",
    "- ```BuyWidgetFromVirtualStore``` adds 1 action set with only 1 action. \n",
    "- ```Gather``` adds 1 action set with 4 actions.\n",
    "- ```ContinuousDoubleAuction``` adds 2\\*N action sets each with M+1 actions, where N is the number of collectible Resources and M is the maximum buying/selling price.\n",
    "\n",
    "In that last example ```ContinuousDoubleAuction``` (which implements trading) the structure of the action sets it creates depend on the choice of maximum buying/selling price (an argument to the class) as well as the collectible Resources in the environment.\n",
    "\n",
    "That might seem complicated but it's simpler than it sounds. Check out [the actual code](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/components/continuous_double_auction.py) for a useful example. In particular, look at the ```__init__``` and ```get_n_agent_actions``` to see how the action sets are set up and look at ```component_step``` to see how the step method makes use of them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Only *you* can ensure NO-OP semantics.\n",
    "\n",
    "Smokey the Bear famously said, \"Only *you* can prevent forest fires.\" While that hazard doesn't apply here, the sense of responsibility is just the same:  \n",
    "**If ```agent.get_component_action(...)``` returns 0, that means NO-OP!**\n",
    "\n",
    "If you look through [the implemented Components](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/components), you'll notice throughout the ```component_step``` methods something along the lines of:\n",
    "```python\n",
    "for agent in self.world.agents:\n",
    "    action = agent.get_component_action(self.name)\n",
    "\n",
    "    # NO-OP! Agent is NOT interacting with this component.\n",
    "    if action == 0: \n",
    "        continue # Move on to the next agent\n",
    "\n",
    "    # Agent is interacting with this component\n",
    "    else:\n",
    "        ... # Do something\n",
    "```\n",
    "\n",
    "Referring back to our ```BuyWidgetFromVirtualStore``` example, notice how it adds 1 action for mobile agents. When an agent actually takes that action we would see ```agent.get_component_action('BuyWidgetFromVirtualStore')``` returns 1. Not 0! If the agent chose an action belonging to another component (say, took a movement action), then we would see ```agent.get_component_action('BuyWidgetFromVirtualStore')``` returns 0, and it would be up to ```BuyWidgetFromVirtualStore``` to treat that like the NO-OP that it is.\n",
    "\n",
    "**When you implement a new Component class, it is up to you to ensure that NO-OP semantics are preserved!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environments come with a couple tools for logging.\n",
    "\n",
    "There are 2 main types of logs that ```BaseEnvironment``` supports: metrics and dense logs.\n",
    "\n",
    "**Metrics** are used to summarize an episode. Scenarios and Components each have a method for producing a metrics dictionary, which adds a tool to generate a readout on what happened. At the end of the episode, any such metrics are combined into a single metrics dictionary which is accessible through ```env.metrics```.\n",
    "\n",
    "**Dense logs** offer a timestep-by-timestep breakdown of how an episode played out. By default, ```BaseEnvironment``` includes world state, agent state, and action info for each timestep. Components can contribute their own dense logs, which get added to the final log at the end of the episode.\n",
    "\n",
    "Because they can be time consuming to create, dense logs are not (by default) generated during every episode. You can use the ```BaseEnvironment``` argument ```dense_log_frequency``` to set how often they are created. If, for example, you use ```dense_log_frequency=20```, then the environment will create dense logs during episodes where the number of total episode completions is a multiple of 20 (that is, every 20 episodes). If you don't want to wait, you can use ```env.reset(force_dense_logging=True)``` to tell the environment to create a dense log for the upcoming episode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Have fun!\n",
    "And congratulate yourself on making it to the end of the advanced tutorial :)\n",
    "\n",
    "If you really want to go for the extra credit, check out [optimal_taxation_theory_and_simulation.ipynb](https://github.com/salesforce/ai-economist/blob/master/tutorials/optimal_taxation_theory_and_simulation.ipynb), our final tutorial which walks through how Foundation is used to study the problem of optimal taxation!"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "948099c6ab02a15a055545cbca87e716aee9b3e6a51d0f68b03005577a92a5b6"
  },
  "kernelspec": {
   "display_name": "Python 3.7.11 64-bit ('ai-economist': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}