tutorials/economic_simulation_basic.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_basic.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Welcome to Foundation!\n",
    "\n",
    "Foundation is the name of the economic simulator built for the AI Economist ([paper here](https://arxiv.org/abs/2004.13332)). This is the first of several tutorials designed to explain how Foundation works and how it can be used to build simulation environments for studying economic problems.\n",
    "\n",
    "Just to orient you a bit, Foundation is specially designed for modeling economies in spatial, 2D grid worlds. The AI Economist paper uses a scenario with 4 agents in a world with *Stone* and *Wood*, which can be *collected*, *traded*, and used to build *Houses*. Here's a (nicely rendered) example of what such an environment looks like:\n",
    "\n",
    "![Foundation snapshot](assets/foundation_snapshot_rendered.jpg)\n",
    "\n",
    "This image just shows what you might see spatially. Behind the scenes, agents have inventories of Stone, Wood, and *Coin*, which they can exchange through a commodities marketplace. In addition, they periodically pay taxes on income earned through trading and building.\n",
    "\n",
    "**We've open-sourced Foundation to foster transparency and to enable others to build on it!** With that goal in mind, this first tutorial should give you enough to see how to create the type of simulation environment described above and how to interact with it. If you're interested to learn how it all works and how to build on it, make sure to check out the advanced tutorial as well! If, after that, you want to understand more about designing the simulation around economic problems, check out our tutorial explaining how the AI Economist uses Foundation to study the optimal taxation problem!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "In this **basic** tutorial, we will demonstrate the basics of how to create an instance of a simulation environment and how to interact with it.\n",
    "\n",
    "We will cover the following: \n",
    "1. Markov Decision Processes\n",
    "2. Creating a Simulation Environment (a Scenario Instance)\n",
    "3. Interacting with the Simulation\n",
    "4. Sampling and Visualizing an Episode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies:\n",
    "You can install the ai-economist package using \n",
    "- the pip package manager OR\n",
    "- by cloning the ai-economist package and installing the requirements (we shall use this when running on Colab):"
   ]
  },
  {
   "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",
    "    !git clone https://github.com/salesforce/ai-economist.git\n",
    "\n",
    "    %cd ai-economist\n",
    "    !pip install -e .\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": [
    "# Import foundation\n",
    "from ai_economist import foundation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change directory to the tutorials folder\n",
    "import os, sys\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "\n",
    "if IN_COLAB:\n",
    "    os.chdir(\"/content/ai-economist/tutorials\")\n",
    "else:\n",
    "    os.chdir(os.path.dirname(os.path.abspath(\"__file__\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "from utils import plotting  # plotting utilities for visualizing env. state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Markov Decision Process\n",
    "\n",
    "Formally, our economic simulation is a key part of a Markov Decision Process (MDP).\n",
    "\n",
    "MDPs describe episodes in which agents interact with a stateful environment in a continuous feedback loop. At each timestep, agents receive an observation and use a policy to choose actions. The environment then advances to a new state, using the old state and the chosen actions. The agents then receive new observations and rewards. This process repeats over $T$ timesteps (possibly infinite).\n",
    "\n",
    "The goal of each agent is to maximize its expected sum of future (discounted) rewards, by finding its optimal policy. Intuitively, this means that an agent needs to understand which (sequence of) actions lead to high rewards (in expectation).\n",
    "\n",
    "### References\n",
    "\n",
    "For more information on reinforcement learning and MDPs, check out:\n",
    "\n",
    "- Richard S. Sutton and Andrew G. Barto, Reinforcement Learning: An Introduction. [http://incompleteideas.net/book/bookdraft2017nov5.pdf](http://incompleteideas.net/book/bookdraft2017nov5.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Creating a Simulation Environment (a Scenario Instance)\n",
    "\n",
    "The Scenario class implements an economic simulation with multiple agents and (optionally) a social planner. \n",
    "\n",
    "Scenarios provide a high-level gym-style API that lets agents interact with it. The API supports multi-agent observations, actions, etc.\n",
    "\n",
    "Each Scenario is stateful and implements two main methods:\n",
    "\n",
    "- __step__, which advances the simulation to the next state, and \n",
    "- __reset__, which puts the simulation back in an initial state.\n",
    "\n",
    "Each Scenario is customizable: you can specify options in a dictionary. Here is an example for a scenario with 4 agents:\n",
    "\n",
    "**Note: This config dictionary will likely seem fairly incomprehensible at this point in the tutorials. Don't worry. The advanced tutorial offers much more context. This is just to get things started and to provide a reference for how to create a \"free market\" economy from the AI Economist.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the configuration of the environment that will be built\n",
    "\n",
    "env_config = {\n",
    "    # ===== SCENARIO CLASS =====\n",
    "    # Which Scenario class to use: the class's name in the Scenario Registry (foundation.scenarios).\n",
    "    # The environment object will be an instance of the Scenario class.\n",
    "    'scenario_name': 'layout_from_file/simple_wood_and_stone',\n",
    "    \n",
    "    # ===== COMPONENTS =====\n",
    "    # Which components to use (specified as list of (\"component_name\", {component_kwargs}) tuples).\n",
    "    #   \"component_name\" refers to the Component class's name in the Component Registry (foundation.components)\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', {'skill_dist': \"pareto\", 'payment_max_skill_multiplier': 3}),\n",
    "        # (2) Trading collectible resources\n",
    "        ('ContinuousDoubleAuction', {'max_num_orders': 5}),\n",
    "        # (3) Movement and resource collection\n",
    "        ('Gather', {}),\n",
    "    ],\n",
    "    \n",
    "    # ===== SCENARIO CLASS ARGUMENTS =====\n",
    "    # (optional) kwargs that are added by the Scenario class (i.e. not defined in BaseEnvironment)\n",
    "    'env_layout_file': 'quadrant_25x25_20each_30clump.txt',\n",
    "    'starting_agent_coin': 10,\n",
    "    'fixed_four_skill_and_loc': True,\n",
    "    \n",
    "    # ===== STANDARD ARGUMENTS ======\n",
    "    # kwargs that are used by every Scenario class (i.e. defined in BaseEnvironment)\n",
    "    'n_agents': 4,          # Number of non-planner agents (must be > 1)\n",
    "    'world_size': [25, 25], # [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 required for masking action logits in the code below.\n",
    "    'flatten_masks': True,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an environment instance using this configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = foundation.make_env_instance(**env_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Interacting with the Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agents\n",
    "\n",
    "The Agent class holds the state of agents in the simulation. Each Agent instance represents a _logical_ agent.\n",
    "\n",
    "_Note that this might be separate from a Policy model that lives outside the Scenario and controls the Agent's behavior._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.get_agent(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A random policy\n",
    "\n",
    "Now let's interact with the simulation.\n",
    "\n",
    "Each Agent needs to choose which actions to execute using a __policy__.\n",
    "\n",
    "Agents might not always be allowed to execute all actions. For instance, a mobile Agent cannot move beyond the boundary of the world. Hence, in position (0, 0), a mobile cannot move \"Left\" or \"Down\". This information is given by a mask, which is provided under ```obs[<agent_id_str>][\"action_mask\"]``` in the observation dictionary ```obs``` returned by the scenario.\n",
    "\n",
    "Let's use a random policy to step through the simulation. The methods below implement a random policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: The code for sampling actions (this cell), and playing an episode (below) are general.  \n",
    "# That is, it doesn't depend on the Scenario and Component classes used in the environment!\n",
    "\n",
    "def sample_random_action(agent, mask):\n",
    "    \"\"\"Sample random UNMASKED action(s) for agent.\"\"\"\n",
    "    # Return a list of actions: 1 for each action subspace\n",
    "    if agent.multi_action_mode:\n",
    "        split_masks = np.split(mask, agent.action_spaces.cumsum()[:-1])\n",
    "        return [np.random.choice(np.arange(len(m_)), p=m_/m_.sum()) for m_ in split_masks]\n",
    "\n",
    "    # Return a single action\n",
    "    else:\n",
    "        return np.random.choice(np.arange(agent.action_spaces), p=mask/mask.sum())\n",
    "\n",
    "def sample_random_actions(env, obs):\n",
    "    \"\"\"Samples random UNMASKED actions for each agent in obs.\"\"\"\n",
    "        \n",
    "    actions = {\n",
    "        a_idx: sample_random_action(env.get_agent(a_idx), a_obs['action_mask'])\n",
    "        for a_idx, a_obs in obs.items()\n",
    "    }\n",
    "\n",
    "    return actions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to interact with the simulation...\n",
    "\n",
    "First, environments can be put in an initial state by using __reset__."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = env.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we call __step__ to advance the state and advance time by one tick."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actions = sample_random_actions(env, obs)\n",
    "obs, rew, done, info = env.step(actions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Internally, the __step__ method composes several Components (which act almost like modular sub-Environments) that implement various agent affordances and environment dynamics. For more detailed information on Components and how to implement custom Component classes, see the advanced tutorial. For this tutorial, we will continue to inspect the information that __step__ returns and run a full episode in the simulation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Observation\n",
    "\n",
    "Each observation is a dictionary that contains information for the $N$ agents and (optionally) social planner (with id \"p\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each agent, the agent-specific observation is a dictionary. Each Component can contribute information to the agent-specific observation. For instance, the Build Component contributes the \n",
    "\n",
    "- Build-build_payment (float)\n",
    "- Build-build_skill (int)\n",
    "\n",
    "fields, which are defined in the ```generate_observations``` method in [foundation/components/build.py](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/components/build.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for key, val in obs['0'].items(): \n",
    "    print(\"{:50} {}\".format(key, type(val)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reward\n",
    "\n",
    "For each agent / planner, the reward dictionary contains a scalar reward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for agent_idx, reward in rew.items(): \n",
    "    print(\"{:2} {:.3f}\".format(agent_idx, reward))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Done\n",
    "\n",
    "The __done__ object is a dictionary that by default records whether all agents have seen the end of the episode. The default criterion for each agent is to 'stop' their episode once $H$ steps have been executed. Once an agent is 'done', they do not change their state anymore. So, while it's not currently implemented, this could be used to indicate that the episode has ended *for a specific Agent*.\n",
    "\n",
    "In general, this is useful for telling a Reinforcement Learning framework when to reset the environment and how to organize the trajectories of individual Agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "done"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Info\n",
    "\n",
    "The __info__ object can record any auxiliary information from the simulator, which can be useful, e.g., for visualization. By default, this is empty. To modify this behavior, modify the step() method in [foundation/base/base_env.py](https://github.com/salesforce/ai-economist/blob/master/ai_economist/foundation/base/base_env.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Sampling and Visualizing an Episode\n",
    "\n",
    "Let's step multiple times with this random policy and visualize the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_plot(env, ax, fig):\n",
    "    \"\"\"Plots world state during episode sampling.\"\"\"\n",
    "    plotting.plot_env_state(env, ax)\n",
    "    ax.set_aspect('equal')\n",
    "    display.display(fig)\n",
    "    display.clear_output(wait=True)\n",
    "\n",
    "def play_random_episode(env, plot_every=100, do_dense_logging=False):\n",
    "    \"\"\"Plays an episode with randomly sampled actions.\n",
    "    \n",
    "    Demonstrates gym-style API:\n",
    "        obs                  <-- env.reset(...)         # Reset\n",
    "        obs, rew, done, info <-- env.step(actions, ...) # Interaction loop\n",
    "    \n",
    "    \"\"\"\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n",
    "\n",
    "    # Reset\n",
    "    obs = env.reset(force_dense_logging=do_dense_logging)\n",
    "\n",
    "    # Interaction loop (w/ plotting)\n",
    "    for t in range(env.episode_length):\n",
    "        actions = sample_random_actions(env, obs)\n",
    "        obs, rew, done, info = env.step(actions)\n",
    "\n",
    "        if ((t+1) % plot_every) == 0:\n",
    "            do_plot(env, ax, fig)\n",
    "\n",
    "    if ((t+1) % plot_every) != 0:\n",
    "        do_plot(env, ax, fig)        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "play_random_episode(env, plot_every=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see four agents (indicated by a circled __\\*__) that move around in the 2-dimensional world. Light brown cells contain Stone, green cells contain Wood. Each agent can build Houses, indicated by corresponding colored cells. Water tiles (blue squares), which prevent movement, divide the map into four quadrants.\n",
    "\n",
    "Note: this is showing the same information as the image at the top of the tutorial -- it just uses a much more simplistic rendering."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize using dense logging\n",
    "\n",
    "Environments built with Foundation provide a couple tools for logging. Perhaps the most useful are **dense logs**. When you reset the environment, you can tell it to create a dense log for the new episode. This will store Agent states at each point in time along with any Component-specific dense log information (say, about builds, trades, etc.) that the Components provide. In addition, it will periodically store a snapshot of the world state.\n",
    "\n",
    "We provide a few plotting tools that work well with the type of environment showcased here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Play another episode. This time, tell the environment to do dense logging\n",
    "play_random_episode(env, plot_every=100, do_dense_logging=True)\n",
    "\n",
    "# Grab the dense log from the env\n",
    "dense_log = env.previous_episode_dense_log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the evolution of the world state from t=0 to t=200\n",
    "fig = plotting.vis_world_range(dense_log, t0=0, tN=200, N=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the evolution of the world state over the full episode\n",
    "fig = plotting.vis_world_range(dense_log, N=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the \"breakdown\" tool to visualize the world state, agent-wise quantities, movement, and trading events\n",
    "plotting.breakdown(dense_log);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Next\n",
    "\n",
    "Now that you've seen how to interact with the simulation and generate episodes, try the next tutorial ([economic_simulation_advanced.ipynb](https://github.com/salesforce/ai-economist/blob/master/tutorials/economic_simulation_advanced.ipynb)) that explains how the simulation is composed of low-level Components and Entities. This structure enables flexible extensions of the economic simulation."
   ]
  }
 ],
 "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
}