adding ai_economist for modding

ai_economist/foundation/components/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2020, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root
+# or https://opensource.org/licenses/BSD-3-Clause
+
+from ai_economist.foundation.base.base_component import component_registry
+
+from . import (
+    build,
+    continuous_double_auction,
+    covid19_components,
+    move,
+    redistribution,
+    simple_labor,
+)
+
+# Import files that add Component class(es) to component_registry
+# ---------------------------------------------------------------

ai_economist/foundation/components/build.py

@@ -0,0 +1,266 @@
+# Copyright (c) 2020, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root
+# or https://opensource.org/licenses/BSD-3-Clause
+
+import numpy as np
+
+from ai_economist.foundation.base.base_component import (
+    BaseComponent,
+    component_registry,
+)
+
+
+@component_registry.add
+class Build(BaseComponent):
+    """
+    Allows mobile agents to build house landmarks in the world using stone and wood,
+    earning income.
+
+    Can be configured to include heterogeneous building skill where agents earn
+    different levels of income when building.
+
+    Args:
+        payment (int): Default amount of coin agents earn from building.
+            Must be >= 0. Default is 10.
+        payment_max_skill_multiplier (int): Maximum skill multiplier that an agent
+            can sample. Must be >= 1. Default is 1.
+        skill_dist (str): Distribution type for sampling skills. Default ("none")
+            gives all agents identical skill equal to a multiplier of 1. "pareto" and
+            "lognormal" sample skills from the associated distributions.
+        build_labor (float): Labor cost associated with building a house.
+            Must be >= 0. Default is 10.
+    """
+
+    name = "Build"
+    component_type = "Build"
+    required_entities = ["Wood", "Stone", "Coin", "House", "Labor"]
+    agent_subclasses = ["BasicMobileAgent"]
+
+    def __init__(
+        self,
+        *base_component_args,
+        payment=10,
+        payment_max_skill_multiplier=1,
+        skill_dist="none",
+        build_labor=10.0,
+        **base_component_kwargs
+    ):
+        super().__init__(*base_component_args, **base_component_kwargs)
+
+        self.payment = int(payment)
+        assert self.payment >= 0
+
+        self.payment_max_skill_multiplier = int(payment_max_skill_multiplier)
+        assert self.payment_max_skill_multiplier >= 1
+
+        self.resource_cost = {"Wood": 1, "Stone": 1}
+
+        self.build_labor = float(build_labor)
+        assert self.build_labor >= 0
+
+        self.skill_dist = skill_dist.lower()
+        assert self.skill_dist in ["none", "pareto", "lognormal"]
+
+        self.sampled_skills = {}
+
+        self.builds = []
+
+    def agent_can_build(self, agent):
+        """Return True if agent can actually build in its current location."""
+        # See if the agent has the resources necessary to complete the action
+        for resource, cost in self.resource_cost.items():
+            if agent.state["inventory"][resource] < cost:
+                return False
+
+        # Do nothing if this spot is already occupied by a landmark or resource
+        if self.world.location_resources(*agent.loc):
+            return False
+        if self.world.location_landmarks(*agent.loc):
+            return False
+        # If we made it here, the agent can build.
+        return True
+
+    # Required methods for implementing components
+    # --------------------------------------------
+
+    def get_n_actions(self, agent_cls_name):
+        """
+        See base_component.py for detailed description.
+
+        Add a single action (build) for mobile agents.
+        """
+        # This component adds 1 action that mobile agents can take: build a house
+        if agent_cls_name == "BasicMobileAgent":
+            return 1
+
+        return None
+
+    def get_additional_state_fields(self, agent_cls_name):
+        """
+        See base_component.py for detailed description.
+
+        For mobile agents, add state fields for building skill.
+        """
+        if agent_cls_name not in self.agent_subclasses:
+            return {}
+        if agent_cls_name == "BasicMobileAgent":
+            return {"build_payment": float(self.payment), "build_skill": 1}
+        raise NotImplementedError
+
+    def component_step(self):
+        """
+        See base_component.py for detailed description.
+
+        Convert stone+wood to house+coin for agents that choose to build and can.
+        """
+        world = self.world
+        build = []
+        # Apply any building actions taken by the mobile agents
+        for agent in world.get_random_order_agents():
+
+            action = agent.get_component_action(self.name)
+
+            # This component doesn't apply to this agent!
+            if action is None:
+                continue
+
+            # NO-OP!
+            if action == 0:
+                pass
+
+            # Build! (If you can.)
+            elif action == 1:
+                if self.agent_can_build(agent):
+                    # Remove the resources
+                    for resource, cost in self.resource_cost.items():
+                        agent.state["inventory"][resource] -= cost
+
+                    # Place a house where the agent is standing
+                    loc_r, loc_c = agent.loc
+                    world.create_landmark("House", loc_r, loc_c, agent.idx)
+
+                    # Receive payment for the house
+                    agent.state["inventory"]["Coin"] += agent.state["build_payment"]
+
+                    # Incur the labor cost for building
+                    agent.state["endogenous"]["Labor"] += self.build_labor
+
+                    build.append(
+                        {
+                            "builder": agent.idx,
+                            "loc": np.array(agent.loc),
+                            "income": float(agent.state["build_payment"]),
+                        }
+                    )
+
+            else:
+                raise ValueError
+
+        self.builds.append(build)
+
+    def generate_observations(self):
+        """
+        See base_component.py for detailed description.
+
+        Here, agents observe their build skill. The planner does not observe anything
+        from this component.
+        """
+
+        obs_dict = dict()
+        for agent in self.world.agents:
+            obs_dict[agent.idx] = {
+                "build_payment": agent.state["build_payment"] / self.payment,
+                "build_skill": self.sampled_skills[agent.idx],
+            }
+
+        return obs_dict
+
+    def generate_masks(self, completions=0):
+        """
+        See base_component.py for detailed description.
+
+        Prevent building only if a landmark already occupies the agent's location.
+        """
+
+        masks = {}
+        # Mobile agents' build action is masked if they cannot build with their
+        # current location and/or endowment
+        for agent in self.world.agents:
+            masks[agent.idx] = np.array([self.agent_can_build(agent)])
+
+        return masks
+
+    # For non-required customization
+    # ------------------------------
+
+    def get_metrics(self):
+        """
+        Metrics that capture what happened through this component.
+
+        Returns:
+            metrics (dict): A dictionary of {"metric_name": metric_value},
+                where metric_value is a scalar.
+        """
+        world = self.world
+
+        build_stats = {a.idx: {"n_builds": 0} for a in world.agents}
+        for builds in self.builds:
+            for build in builds:
+                idx = build["builder"]
+                build_stats[idx]["n_builds"] += 1
+
+        out_dict = {}
+        for a in world.agents:
+            for k, v in build_stats[a.idx].items():
+                out_dict["{}/{}".format(a.idx, k)] = v
+
+        num_houses = np.sum(world.maps.get("House") > 0)
+        out_dict["total_builds"] = num_houses
+
+        return out_dict
+
+    def additional_reset_steps(self):
+        """
+        See base_component.py for detailed description.
+
+        Re-sample agents' building skills.
+        """
+        world = self.world
+
+        self.sampled_skills = {agent.idx: 1 for agent in world.agents}
+
+        PMSM = self.payment_max_skill_multiplier
+
+        for agent in world.agents:
+            if self.skill_dist == "none":
+                sampled_skill = 1
+                pay_rate = 1
+            elif self.skill_dist == "pareto":
+                sampled_skill = np.random.pareto(4)
+                pay_rate = np.minimum(PMSM, (PMSM - 1) * sampled_skill + 1)
+            elif self.skill_dist == "lognormal":
+                sampled_skill = np.random.lognormal(-1, 0.5)
+                pay_rate = np.minimum(PMSM, (PMSM - 1) * sampled_skill + 1)
+            else:
+                raise NotImplementedError
+
+            agent.state["build_payment"] = float(pay_rate * self.payment)
+            agent.state["build_skill"] = float(sampled_skill)
+
+            self.sampled_skills[agent.idx] = sampled_skill
+
+        self.builds = []
+
+    def get_dense_log(self):
+        """
+        Log builds.
+
+        Returns:
+            builds (list): A list of build events. Each entry corresponds to a single
+                timestep and contains a description of any builds that occurred on
+                that timestep.
+
+        """
+        return self.builds

ai_economist/foundation/components/continuous_double_auction.py

@@ -0,0 +1,679 @@
+# Copyright (c) 2020, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root
+# or https://opensource.org/licenses/BSD-3-Clause
+
+import numpy as np
+
+from ai_economist.foundation.base.base_component import (
+    BaseComponent,
+    component_registry,
+)
+from ai_economist.foundation.entities import resource_registry
+
+
+@component_registry.add
+class ContinuousDoubleAuction(BaseComponent):
+    """Allows mobile agents to buy/sell collectible resources with one another.
+
+    Implements a commodity-exchange-style market where agents may sell a unit of
+        resource by submitting an ask (saying the minimum it will accept in payment)
+        or may buy a resource by submitting a bid (saying the maximum it will pay in
+        exchange for a unit of a given resource).
+
+    Args:
+        max_bid_ask (int): Maximum amount of coin that an agent can bid or ask for.
+            Must be >= 1. Default is 10 coin.
+        order_labor (float): Amount of labor incurred when an agent creates an order.
+            Must be >= 0. Default is 0.25.
+        order_duration (int): Number of environment timesteps before an unfilled
+            bid/ask expires. Must be >= 1. Default is 50 timesteps.
+        max_num_orders (int, optional): Maximum number of bids + asks that an agent can
+            have open for a given resource. Must be >= 1. Default is no limit to
+            number of orders.
+    """
+
+    name = "ContinuousDoubleAuction"
+    component_type = "Trade"
+    required_entities = ["Coin", "Labor"]
+    agent_subclasses = ["BasicMobileAgent"]
+
+    def __init__(
+        self,
+        *args,
+        max_bid_ask=10,
+        order_labor=0.25,
+        order_duration=50,
+        max_num_orders=None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+
+        # The max amount (in coin) that an agent can bid/ask for 1 unit of a commodity
+        self.max_bid_ask = int(max_bid_ask)
+        assert self.max_bid_ask >= 1
+        self.price_floor = 0
+        self.price_ceiling = int(max_bid_ask)
+
+        # The amount of time (in timesteps) that an order stays in the books
+        # before it expires
+        self.order_duration = int(order_duration)
+        assert self.order_duration >= 1
+
+        # The maximum number of bid+ask orders an agent can have open
+        # for each type of commodity
+        self.max_num_orders = int(max_num_orders or self.order_duration)
+        assert self.max_num_orders >= 1
+
+        # The labor cost associated with creating a bid or ask order
+
+        self.order_labor = float(order_labor)
+        self.order_labor = max(self.order_labor, 0.0)
+
+        # Each collectible resource in the world can be traded via this component
+        self.commodities = [
+            r for r in self.world.resources if resource_registry.get(r).collectible
+        ]
+
+        # These get reset at the start of an episode:
+        self.asks = {c: [] for c in self.commodities}
+        self.bids = {c: [] for c in self.commodities}
+        self.n_orders = {
+            c: {i: 0 for i in range(self.n_agents)} for c in self.commodities
+        }
+        self.executed_trades = []
+        self.price_history = {
+            c: {i: self._price_zeros() for i in range(self.n_agents)}
+            for c in self.commodities
+        }
+        self.bid_hists = {
+            c: {i: self._price_zeros() for i in range(self.n_agents)}
+            for c in self.commodities
+        }
+        self.ask_hists = {
+            c: {i: self._price_zeros() for i in range(self.n_agents)}
+            for c in self.commodities
+        }
+
+    # Convenience methods
+    # -------------------
+
+    def _price_zeros(self):
+        if 1 + self.price_ceiling - self.price_floor <= 0:
+            print("ERROR!", self.price_ceiling, self.price_floor)
+
+        return np.zeros(1 + self.price_ceiling - self.price_floor)
+
+    def available_asks(self, resource, agent):
+        """
+        Get a histogram of asks for resource to which agent could bid against.
+
+        Args:
+            resource (str): Name of the resource
+            agent (BasicMobileAgent or None): Object of agent for which available
+                asks are being queried. If None, all asks are considered available.
+
+        Returns:
+            ask_hist (ndarray): For each possible price level, the number of
+                available asks.
+        """
+        if agent is None:
+            a_idx = -1
+        else:
+            a_idx = agent.idx
+        ask_hist = self._price_zeros()
+        for i, h in self.ask_hists[resource].items():
+            if a_idx != i:
+                ask_hist += h
+        return ask_hist
+
+    def available_bids(self, resource, agent):
+        """
+        Get a histogram of bids for resource to which agent could ask against.
+
+        Args:
+            resource (str): Name of the resource
+            agent (BasicMobileAgent or None): Object of agent for which available
+                bids are being queried. If None, all bids are considered available.
+
+        Returns:
+            bid_hist (ndarray): For each possible price level, the number of
+                available bids.
+        """
+        if agent is None:
+            a_idx = -1
+        else:
+            a_idx = agent.idx
+        bid_hist = self._price_zeros()
+        for i, h in self.bid_hists[resource].items():
+            if a_idx != i:
+                bid_hist += h
+        return bid_hist
+
+    def can_bid(self, resource, agent):
+        """If agent can submit a bid for resource."""
+        return self.n_orders[resource][agent.idx] < self.max_num_orders
+
+    def can_ask(self, resource, agent):
+        """If agent can submit an ask for resource."""
+        return (
+            self.n_orders[resource][agent.idx] < self.max_num_orders
+            and agent.state["inventory"][resource] > 0
+        )
+
+    # Core components for this market
+    # -------------------------------
+
+    def create_bid(self, resource, agent, max_payment):
+        """Create a new bid for resource, with agent offering max_payment.
+
+        On a successful trade, payment will be at most max_payment, possibly less.
+
+        The agent places the bid coin into escrow so that it may not be spent on
+        something else while the order exists.
+        """
+
+        # The agent is past the max number of orders
+        # or doesn't have enough money, do nothing
+        if (not self.can_bid(resource, agent)) or agent.state["inventory"][
+            "Coin"
+        ] < max_payment:
+            return
+
+        assert self.price_floor <= max_payment <= self.price_ceiling
+
+        bid = {"buyer": agent.idx, "bid": int(max_payment), "bid_lifetime": 0}
+
+        # Add this to the bid book
+        self.bids[resource].append(bid)
+        self.bid_hists[resource][bid["buyer"]][bid["bid"] - self.price_floor] += 1
+        self.n_orders[resource][agent.idx] += 1
+
+        # Set aside whatever money the agent is willing to pay
+        # (will get excess back if price ends up being less)
+        _ = agent.inventory_to_escrow("Coin", int(max_payment))
+
+        # Incur the labor cost of creating an order
+        agent.state["endogenous"]["Labor"] += self.order_labor
+
+    def create_ask(self, resource, agent, min_income):
+        """
+        Create a new ask for resource, with agent asking for min_income.
+
+        On a successful trade, income will be at least min_income, possibly more.
+
+        The agent places one unit of resource into escrow so that it may not be used
+        for something else while the order exists.
+        """
+        # The agent is past the max number of orders
+        # or doesn't the resource it's trying to sell, do nothing
+        if not self.can_ask(resource, agent):
+            return
+
+        # is there an upper limit?
+        assert self.price_floor <= min_income <= self.price_ceiling
+
+        ask = {"seller": agent.idx, "ask": int(min_income), "ask_lifetime": 0}
+
+        # Add this to the ask book
+        self.asks[resource].append(ask)
+        self.ask_hists[resource][ask["seller"]][ask["ask"] - self.price_floor] += 1
+        self.n_orders[resource][agent.idx] += 1
+
+        # Set aside the resource the agent is willing to sell
+        amount = agent.inventory_to_escrow(resource, 1)
+        assert amount == 1
+
+        # Incur the labor cost of creating an order
+        agent.state["endogenous"]["Labor"] += self.order_labor
+
+    def match_orders(self):
+        """
+        This implements the continuous double auction by identifying valid bid/ask
+        pairs and executing trades accordingly.
+
+        Higher (lower) bids (asks) are given priority over lower (higher) bids (asks).
+        Trades are executed using the price of whichever bid/ask order was placed
+        first: bid price if bid was placed first, ask price otherwise.
+
+        Trading removes the payment and resource from bidder's and asker's escrow,
+        respectively, and puts them in the other's inventory.
+        """
+        self.executed_trades.append([])
+
+        for resource in self.commodities:
+            possible_match = [True for _ in range(self.n_agents)]
+            keep_checking = True
+
+            bids = sorted(
+                self.bids[resource],
+                key=lambda b: (b["bid"], b["bid_lifetime"]),
+                reverse=True,
+            )
+            asks = sorted(
+                self.asks[resource], key=lambda a: (a["ask"], -a["ask_lifetime"])
+            )
+
+            while any(possible_match) and keep_checking:
+                idx_bid, idx_ask = 0, 0
+                while True:
+                    # Out of bids to check. Exit both loops.
+                    if idx_bid >= len(bids):
+                        keep_checking = False
+                        break
+
+                    # Already know this buyer is no good for this round.
+                    # Skip to next bid.
+                    if not possible_match[bids[idx_bid]["buyer"]]:
+                        idx_bid += 1
+
+                    # Out of asks to check. This buyer won't find a match on this round.
+                    # (maybe) Restart inner loop.
+                    elif idx_ask >= len(asks):
+                        possible_match[bids[idx_bid]["buyer"]] = False
+                        break
+
+                    # Skip to next ask if this ask comes from the buyer
+                    # of the current bid.
+                    elif asks[idx_ask]["seller"] == bids[idx_bid]["buyer"]:
+                        idx_ask += 1
+
+                    # If this bid/ask pair can't be matched, this buyer
+                    # can't be matched. (maybe) Restart inner loop.
+                    elif bids[idx_bid]["bid"] < asks[idx_ask]["ask"]:
+                        possible_match[bids[idx_bid]["buyer"]] = False
+                        break
+
+                    # TRADE! (then restart inner loop)
+                    else:
+                        bid = bids.pop(idx_bid)
+                        ask = asks.pop(idx_ask)
+
+                        trade = {"commodity": resource}
+                        trade.update(bid)
+                        trade.update(ask)
+
+                        if (
+                            bid["bid_lifetime"] <= ask["ask_lifetime"]
+                        ):  # Ask came earlier. (in other words,
+                            # trade triggered by new bid)
+                            trade["price"] = int(trade["ask"])
+                        else:  # Bid came earlier. (in other words,
+                            # trade triggered by new ask)
+                            trade["price"] = int(trade["bid"])
+                        trade["cost"] = trade["price"]  # What the buyer pays in total
+                        trade["income"] = trade[
+                            "price"
+                        ]  # What the seller receives in total
+
+                        buyer = self.world.agents[trade["buyer"]]
+                        seller = self.world.agents[trade["seller"]]
+
+                        # Bookkeeping
+                        self.bid_hists[resource][bid["buyer"]][
+                            bid["bid"] - self.price_floor
+                        ] -= 1
+                        self.ask_hists[resource][ask["seller"]][
+                            ask["ask"] - self.price_floor
+                        ] -= 1
+                        self.n_orders[trade["commodity"]][seller.idx] -= 1
+                        self.n_orders[trade["commodity"]][buyer.idx] -= 1
+                        self.executed_trades[-1].append(trade)
+                        self.price_history[resource][trade["seller"]][
+                            trade["price"]
+                        ] += 1
+
+                        # The resource goes from the seller's escrow
+                        # to the buyer's inventory
+                        seller.state["escrow"][resource] -= 1
+                        buyer.state["inventory"][resource] += 1
+
+                        # Buyer's money (already set aside) leaves escrow
+                        pre_payment = int(trade["bid"])
+                        buyer.state["escrow"]["Coin"] -= pre_payment
+                        assert buyer.state["escrow"]["Coin"] >= 0
+
+                        # Payment is removed from the pre_payment
+                        # and given to the seller. Excess returned to buyer.
+                        payment_to_seller = int(trade["price"])
+                        excess_payment_from_buyer = pre_payment - payment_to_seller
+                        assert excess_payment_from_buyer >= 0
+                        seller.state["inventory"]["Coin"] += payment_to_seller
+                        buyer.state["inventory"]["Coin"] += excess_payment_from_buyer
+
+                        # Restart the inner loop
+                        break
+
+            # Keep the unfilled bids/asks
+            self.bids[resource] = bids
+            self.asks[resource] = asks
+
+    def remove_expired_orders(self):
+        """
+        Increment the time counter for any unfilled bids/asks and remove expired
+        orders from the market.
+
+        When orders expire, the payment or resource is removed from escrow and
+        returned to the inventory and the associated order is removed from the order
+        books.
+        """
+        world = self.world
+
+        for resource in self.commodities:
+
+            bids_ = []
+            for bid in self.bids[resource]:
+                bid["bid_lifetime"] += 1
+                # If the bid is not expired, keep it in the bids
+                if bid["bid_lifetime"] <= self.order_duration:
+                    bids_.append(bid)
+                # Otherwise, remove it and do the associated bookkeeping
+                else:
+                    # Return the set aside money to the buyer
+                    amount = world.agents[bid["buyer"]].escrow_to_inventory(
+                        "Coin", bid["bid"]
+                    )
+                    assert amount == bid["bid"]
+                    # Adjust the bid histogram to reflect the removal of the bid
+                    self.bid_hists[resource][bid["buyer"]][
+                        bid["bid"] - self.price_floor
+                    ] -= 1
+                    # Adjust the order counter
+                    self.n_orders[resource][bid["buyer"]] -= 1
+
+            asks_ = []
+            for ask in self.asks[resource]:
+                ask["ask_lifetime"] += 1
+                # If the ask is not expired, keep it in the asks
+                if ask["ask_lifetime"] <= self.order_duration:
+                    asks_.append(ask)
+                # Otherwise, remove it and do the associated bookkeeping
+                else:
+                    # Return the set aside resource to the seller
+                    resource_unit = world.agents[ask["seller"]].escrow_to_inventory(
+                        resource, 1
+                    )
+                    assert resource_unit == 1
+                    # Adjust the ask histogram to reflect the removal of the ask
+                    self.ask_hists[resource][ask["seller"]][
+                        ask["ask"] - self.price_floor
+                    ] -= 1
+                    # Adjust the order counter
+                    self.n_orders[resource][ask["seller"]] -= 1
+
+            self.bids[resource] = bids_
+            self.asks[resource] = asks_
+
+    # Required methods for implementing components
+    # --------------------------------------------
+
+    def get_n_actions(self, agent_cls_name):
+        """
+        See base_component.py for detailed description.
+
+        Adds 2*C action spaces [ (bid+ask) * n_commodities ], each with 1 + max_bid_ask
+        actions corresponding to price levels 0 to max_bid_ask.
+        """
+        # This component adds 2*(1+max_bid_ask)*n_resources possible actions:
+        # buy/sell x each-price x each-resource
+        if agent_cls_name == "BasicMobileAgent":
+            trades = []
+            for c in self.commodities:
+                trades.append(
+                    ("Buy_{}".format(c), 1 + self.max_bid_ask)
+                )  # How much willing to pay for c
+                trades.append(
+                    ("Sell_{}".format(c), 1 + self.max_bid_ask)
+                )  # How much need to receive to sell c
+            return trades
+
+        return None
+
+    def get_additional_state_fields(self, agent_cls_name):
+        """
+        See base_component.py for detailed description.
+        """
+        # This component doesn't add any state fields
+        return {}
+
+    def component_step(self):
+        """
+        See base_component.py for detailed description.
+
+        Create new bids and asks, match and execute valid order pairs, and manage
+        order expiration.
+        """
+        world = self.world
+
+        for resource in self.commodities:
+            for agent in world.agents:
+                self.price_history[resource][agent.idx] *= 0.995
+
+                # Create bid action
+                # -----------------
+                resource_action = agent.get_component_action(
+                    self.name, "Buy_{}".format(resource)
+                )
+
+                # No-op
+                if resource_action == 0:
+                    pass
+
+                # Create a bid
+                elif resource_action <= self.max_bid_ask + 1:
+                    self.create_bid(resource, agent, max_payment=resource_action - 1)
+
+                else:
+                    raise ValueError
+
+                # Create ask action
+                # -----------------
+                resource_action = agent.get_component_action(
+                    self.name, "Sell_{}".format(resource)
+                )
+
+                # No-op
+                if resource_action == 0:
+                    pass
+
+                # Create an ask
+                elif resource_action <= self.max_bid_ask + 1:
+                    self.create_ask(resource, agent, min_income=resource_action - 1)
+
+                else:
+                    raise ValueError
+
+        # Here's where the magic happens:
+        self.match_orders()  # Pair bids and asks
+        self.remove_expired_orders()  # Get rid of orders that have expired
+
+    def generate_observations(self):
+        """
+        See base_component.py for detailed description.
+
+        Here, agents and the planner both observe historical market behavior and
+        outstanding bids/asks for each tradable commodity. Agents only see the
+        outstanding bids/asks to which they could respond (that is, that they did not
+        submit). Agents also see their own outstanding bids/asks.
+        """
+        world = self.world
+
+        obs = {a.idx: {} for a in world.agents + [world.planner]}
+
+        prices = np.arange(self.price_floor, self.price_ceiling + 1)
+        for c in self.commodities:
+            net_price_history = np.sum(
+                np.stack([self.price_history[c][i] for i in range(self.n_agents)]),
+                axis=0,
+            )
+            market_rate = prices.dot(net_price_history) / np.maximum(
+                0.001, np.sum(net_price_history)
+            )
+            scaled_price_history = net_price_history * self.inv_scale
+
+            full_asks = self.available_asks(c, agent=None)
+            full_bids = self.available_bids(c, agent=None)
+
+            obs[world.planner.idx].update(
+                {
+                    "market_rate-{}".format(c): market_rate,
+                    "price_history-{}".format(c): scaled_price_history,
+                    "full_asks-{}".format(c): full_asks,
+                    "full_bids-{}".format(c): full_bids,
+                }
+            )
+
+            for _, agent in enumerate(world.agents):
+                # Private to the agent
+                obs[agent.idx].update(
+                    {
+                        "market_rate-{}".format(c): market_rate,
+                        "price_history-{}".format(c): scaled_price_history,
+                        "available_asks-{}".format(c): full_asks
+                        - self.ask_hists[c][agent.idx],
+                        "available_bids-{}".format(c): full_bids
+                        - self.bid_hists[c][agent.idx],
+                        "my_asks-{}".format(c): self.ask_hists[c][agent.idx],
+                        "my_bids-{}".format(c): self.bid_hists[c][agent.idx],
+                    }
+                )
+
+        return obs
+
+    def generate_masks(self, completions=0):
+        """
+        See base_component.py for detailed description.
+
+        Agents cannot submit bids/asks for resources where they are at the order
+        limit. In addition, they may only submit asks for resources they possess and
+        bids for which they can pay.
+        """
+        world = self.world
+
+        masks = dict()
+
+        for agent in world.agents:
+            masks[agent.idx] = {}
+
+            can_pay = np.arange(self.max_bid_ask + 1) <= agent.inventory["Coin"]
+
+            for resource in self.commodities:
+                if not self.can_ask(resource, agent):  # asks_maxed:
+                    masks[agent.idx]["Sell_{}".format(resource)] = np.zeros(
+                        1 + self.max_bid_ask
+                    )
+                else:
+                    masks[agent.idx]["Sell_{}".format(resource)] = np.ones(
+                        1 + self.max_bid_ask
+                    )
+
+                if not self.can_bid(resource, agent):
+                    masks[agent.idx]["Buy_{}".format(resource)] = np.zeros(
+                        1 + self.max_bid_ask
+                    )
+                else:
+                    masks[agent.idx]["Buy_{}".format(resource)] = can_pay.astype(
+                        np.int32
+                    )
+
+        return masks
+
+    # For non-required customization
+    # ------------------------------
+
+    def get_metrics(self):
+        """
+        Metrics that capture what happened through this component.
+
+        Returns:
+            metrics (dict): A dictionary of {"metric_name": metric_value},
+                where metric_value is a scalar.
+        """
+        world = self.world
+
+        trade_keys = ["price", "cost", "income"]
+
+        selling_stats = {
+            a.idx: {
+                c: {k: 0 for k in trade_keys + ["n_sales"]} for c in self.commodities
+            }
+            for a in world.agents
+        }
+        buying_stats = {
+            a.idx: {
+                c: {k: 0 for k in trade_keys + ["n_sales"]} for c in self.commodities
+            }
+            for a in world.agents
+        }
+
+        n_trades = 0
+
+        for trades in self.executed_trades:
+            for trade in trades:
+                n_trades += 1
+                i_s, i_b, c = trade["seller"], trade["buyer"], trade["commodity"]
+                selling_stats[i_s][c]["n_sales"] += 1
+                buying_stats[i_b][c]["n_sales"] += 1
+                for k in trade_keys:
+                    selling_stats[i_s][c][k] += trade[k]
+                    buying_stats[i_b][c][k] += trade[k]
+
+        out_dict = {}
+        for a in world.agents:
+            for c in self.commodities:
+                for stats, prefix in zip(
+                    [selling_stats, buying_stats], ["Sell", "Buy"]
+                ):
+                    n = stats[a.idx][c]["n_sales"]
+                    if n == 0:
+                        for k in trade_keys:
+                            stats[a.idx][c][k] = np.nan
+                    else:
+                        for k in trade_keys:
+                            stats[a.idx][c][k] /= n
+
+                    for k, v in stats[a.idx][c].items():
+                        out_dict["{}/{}{}/{}".format(a.idx, prefix, c, k)] = v
+
+        out_dict["n_trades"] = n_trades
+
+        return out_dict
+
+    def additional_reset_steps(self):
+        """
+        See base_component.py for detailed description.
+
+        Reset the order books.
+        """
+        self.bids = {c: [] for c in self.commodities}
+        self.asks = {c: [] for c in self.commodities}
+        self.n_orders = {
+            c: {i: 0 for i in range(self.n_agents)} for c in self.commodities
+        }
+
+        self.price_history = {
+            c: {i: self._price_zeros() for i in range(self.n_agents)}
+            for c in self.commodities
+        }
+        self.bid_hists = {
+            c: {i: self._price_zeros() for i in range(self.n_agents)}
+            for c in self.commodities
+        }
+        self.ask_hists = {
+            c: {i: self._price_zeros() for i in range(self.n_agents)}
+            for c in self.commodities
+        }
+
+        self.executed_trades = []
+
+    def get_dense_log(self):
+        """
+        Log executed trades.
+
+        Returns:
+            trades (list): A list of trade events. Each entry corresponds to a single
+                timestep and contains a description of any trades that occurred on
+                that timestep.
+        """
+        return self.executed_trades

ai_economist/foundation/components/covid19_components.py

@@ -0,0 +1,663 @@
+# Copyright (c) 2021, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root
+# or https://opensource.org/licenses/BSD-3-Clause
+
+from datetime import datetime
+
+import GPUtil
+import numpy as np
+
+from ai_economist.foundation.base.base_component import (
+    BaseComponent,
+    component_registry,
+)
+
+try:
+    num_gpus_available = len(GPUtil.getAvailable())
+    print(f"Inside covid19_components.py: {num_gpus_available} GPUs are available.")
+    if num_gpus_available == 0:
+        print("No GPUs found! Running the simulation on a CPU.")
+    else:
+        from warp_drive.utils.constants import Constants
+        from warp_drive.utils.data_feed import DataFeed
+
+        _OBSERVATIONS = Constants.OBSERVATIONS
+        _ACTIONS = Constants.ACTIONS
+except ModuleNotFoundError:
+    print(
+        "Warning: The 'WarpDrive' package is not found and cannot be used! "
+        "If you wish to use WarpDrive, please run "
+        "'pip install rl-warp-drive' first."
+    )
+except ValueError:
+    print("No GPUs found! Running the simulation on a CPU.")
+
+
+@component_registry.add
+class ControlUSStateOpenCloseStatus(BaseComponent):
+    """
+    Sets the open/close stringency levels for states.
+    Args:
+        n_stringency_levels (int): number of stringency levels the states can chose
+            from. (Must match the number in the model constants dictionary referenced by
+            the parent scenario.)
+        action_cooldown_period (int): action cooldown period in days.
+            Once a stringency level is set, the state(s) cannot switch to another level
+            for a certain number of days (referred to as the "action_cooldown_period")
+    """
+
+    name = "ControlUSStateOpenCloseStatus"
+    required_entities = []
+    agent_subclasses = ["BasicMobileAgent"]
+
+    def __init__(
+        self,
+        *base_component_args,
+        n_stringency_levels=10,
+        action_cooldown_period=28,
+        **base_component_kwargs,
+    ):
+
+        self.action_cooldown_period = action_cooldown_period
+        super().__init__(*base_component_args, **base_component_kwargs)
+        self.np_int_dtype = np.int32
+
+        self.n_stringency_levels = int(n_stringency_levels)
+        assert self.n_stringency_levels >= 2
+        self._checked_n_stringency_levels = False
+
+        self.masks = dict()
+        self.default_agent_action_mask = [1 for _ in range(self.n_stringency_levels)]
+        self.no_op_agent_action_mask = [0 for _ in range(self.n_stringency_levels)]
+        self.masks["a"] = np.repeat(
+            np.array(self.no_op_agent_action_mask)[:, np.newaxis],
+            self.n_agents,
+            axis=-1,
+        )
+
+        # (This will be overwritten during reset; see below)
+        self.action_in_cooldown_until = None
+
+    def get_additional_state_fields(self, agent_cls_name):
+        return {}
+
+    def additional_reset_steps(self):
+        # Store the times when the next set of actions can be taken.
+        self.action_in_cooldown_until = np.array(
+            [self.world.timestep for _ in range(self.n_agents)]
+        )
+
+    def get_n_actions(self, agent_cls_name):
+        if agent_cls_name == "BasicMobileAgent":
+            return self.n_stringency_levels
+        return None
+
+    def generate_masks(self, completions=0):
+        for agent in self.world.agents:
+            if self.world.use_real_world_policies:
+                self.masks["a"][:, agent.idx] = self.default_agent_action_mask
+            else:
+                if self.world.timestep < self.action_in_cooldown_until[agent.idx]:
+                    # Keep masking the actions
+                    self.masks["a"][:, agent.idx] = self.no_op_agent_action_mask
+                else:  # self.world.timestep == self.action_in_cooldown_until[agent.idx]
+                    # Cooldown period has ended; unmask the "subsequent" action
+                    self.masks["a"][:, agent.idx] = self.default_agent_action_mask
+        return self.masks
+
+    def get_data_dictionary(self):
+        """
+        Create a dictionary of data to push to the GPU (device).
+        """
+        data_dict = DataFeed()
+        data_dict.add_data(
+            name="action_cooldown_period",
+            data=self.action_cooldown_period,
+        )
+        data_dict.add_data(
+            name="action_in_cooldown_until",
+            data=self.action_in_cooldown_until,
+            save_copy_and_apply_at_reset=True,
+        )
+        data_dict.add_data(
+            name="num_stringency_levels",
+            data=self.n_stringency_levels,
+        )
+        data_dict.add_data(
+            name="default_agent_action_mask",
+            data=[1] + self.default_agent_action_mask,
+        )
+        data_dict.add_data(
+            name="no_op_agent_action_mask",
+            data=[1] + self.no_op_agent_action_mask,
+        )
+        return data_dict
+
+    def get_tensor_dictionary(self):
+        """
+        Create a dictionary of (Pytorch-accessible) data to push to the GPU (device).
+        """
+        tensor_dict = DataFeed()
+        return tensor_dict
+
+    def component_step(self):
+        if self.world.use_cuda:
+            self.world.cuda_component_step[self.name](
+                self.world.cuda_data_manager.device_data("stringency_level"),
+                self.world.cuda_data_manager.device_data("action_cooldown_period"),
+                self.world.cuda_data_manager.device_data("action_in_cooldown_until"),
+                self.world.cuda_data_manager.device_data("default_agent_action_mask"),
+                self.world.cuda_data_manager.device_data("no_op_agent_action_mask"),
+                self.world.cuda_data_manager.device_data("num_stringency_levels"),
+                self.world.cuda_data_manager.device_data(f"{_ACTIONS}_a"),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_a_{self.name}-agent_policy_indicators"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_a_action_mask"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_p_{self.name}-agent_policy_indicators"
+                ),
+                self.world.cuda_data_manager.device_data("_timestep_"),
+                self.world.cuda_data_manager.meta_info("n_agents"),
+                self.world.cuda_data_manager.meta_info("episode_length"),
+                block=self.world.cuda_function_manager.block,
+                grid=self.world.cuda_function_manager.grid,
+            )
+        else:
+            if not self._checked_n_stringency_levels:
+                if self.n_stringency_levels != self.world.n_stringency_levels:
+                    raise ValueError(
+                        "The environment was not configured correctly. For the given "
+                        "model fit, you need to set the number of stringency levels to "
+                        "be {}".format(self.world.n_stringency_levels)
+                    )
+                self._checked_n_stringency_levels = True
+
+            for agent in self.world.agents:
+                if self.world.use_real_world_policies:
+                    # Use the action taken in the previous timestep
+                    action = self.world.real_world_stringency_policy[
+                        self.world.timestep - 1, agent.idx
+                    ]
+                else:
+                    action = agent.get_component_action(self.name)
+                assert 0 <= action <= self.n_stringency_levels
+
+                # We only update the stringency level if the action is not a NO-OP.
+                self.world.global_state["Stringency Level"][
+                    self.world.timestep, agent.idx
+                ] = (
+                    self.world.global_state["Stringency Level"][
+                        self.world.timestep - 1, agent.idx
+                    ]
+                    * (action == 0)
+                    + action
+                )
+
+                agent.state[
+                    "Current Open Close Stringency Level"
+                ] = self.world.global_state["Stringency Level"][
+                    self.world.timestep, agent.idx
+                ]
+
+                # Check if the action cooldown period has ended, and set the next
+                # time until action cooldown. If current action is a no-op
+                # (i.e., no new action was taken), the agent can take an action
+                # in the very next step, otherwise it needs to wait for
+                # self.action_cooldown_period steps. When in the action cooldown
+                # period, whatever actions the agents take are masked out,
+                # so it's always a NO-OP (see generate_masks() above)
+                # The logic below influences the action masks.
+                if self.world.timestep == self.action_in_cooldown_until[agent.idx] + 1:
+                    if action == 0:  # NO-OP
+                        self.action_in_cooldown_until[agent.idx] += 1
+                    else:
+                        self.action_in_cooldown_until[
+                            agent.idx
+                        ] += self.action_cooldown_period
+
+    def generate_observations(self):
+
+        # Normalized observations
+        obs_dict = dict()
+        agent_policy_indicators = self.world.global_state["Stringency Level"][
+            self.world.timestep
+        ]
+        obs_dict["a"] = {
+            "agent_policy_indicators": agent_policy_indicators
+            / self.n_stringency_levels
+        }
+        obs_dict[self.world.planner.idx] = {
+            "agent_policy_indicators": agent_policy_indicators
+            / self.n_stringency_levels
+        }
+
+        return obs_dict
+
+
+@component_registry.add
+class FederalGovernmentSubsidy(BaseComponent):
+    """
+    Args:
+        subsidy_interval (int): The number of days over which the total subsidy amount
+            is evenly rolled out.
+            Note: shortening the subsidy interval increases the total amount of money
+            that the planner could possibly spend. For instance, if the subsidy
+            interval is 30, the planner can create a subsidy every 30 days.
+        num_subsidy_levels (int): The number of subsidy levels.
+            Note: with max_annual_subsidy_per_person=10000, one round of subsidies at
+            the maximum subsidy level equals an expenditure of roughly $3.3 trillion
+            (given the US population of 330 million).
+            If the planner chooses the maximum subsidy amount, the $3.3 trillion
+            is rolled out gradually over the subsidy interval.
+        max_annual_subsidy_per_person (float): The maximum annual subsidy that may be
+            allocated per person.
+    """
+
+    name = "FederalGovernmentSubsidy"
+    required_entities = []
+    agent_subclasses = ["BasicPlanner"]
+
+    def __init__(
+        self,
+        *base_component_args,
+        subsidy_interval=90,
+        num_subsidy_levels=20,
+        max_annual_subsidy_per_person=20000,
+        **base_component_kwargs,
+    ):
+        self.subsidy_interval = int(subsidy_interval)
+        assert self.subsidy_interval >= 1
+
+        self.num_subsidy_levels = int(num_subsidy_levels)
+        assert self.num_subsidy_levels >= 1
+
+        self.max_annual_subsidy_per_person = float(max_annual_subsidy_per_person)
+        assert self.max_annual_subsidy_per_person >= 0
+
+        self.np_int_dtype = np.int32
+
+        # (This will be overwritten during component_step; see below)
+        self._subsidy_amount_per_level = None
+        self._subsidy_level_array = None
+
+        super().__init__(*base_component_args, **base_component_kwargs)
+
+        self.default_planner_action_mask = [1 for _ in range(self.num_subsidy_levels)]
+        self.no_op_planner_action_mask = [0 for _ in range(self.num_subsidy_levels)]
+
+        # (This will be overwritten during reset; see below)
+        self.max_daily_subsidy_per_state = np.array(
+            self.n_agents, dtype=self.np_int_dtype
+        )
+
+    def get_additional_state_fields(self, agent_cls_name):
+        if agent_cls_name == "BasicPlanner":
+            return {"Total Subsidy": 0, "Current Subsidy Level": 0}
+        return {}
+
+    def additional_reset_steps(self):
+        # Pre-compute maximum state-specific subsidy levels
+        self.max_daily_subsidy_per_state = (
+            self.world.us_state_population * self.max_annual_subsidy_per_person / 365
+        )
+
+    def get_n_actions(self, agent_cls_name):
+        if agent_cls_name == "BasicPlanner":
+            # Number of non-zero subsidy levels
+            # (the action 0 pertains to the no-subsidy case)
+            return self.num_subsidy_levels
+        return None
+
+    def generate_masks(self, completions=0):
+        masks = {}
+        if self.world.use_real_world_policies:
+            masks[self.world.planner.idx] = self.default_planner_action_mask
+        else:
+            if self.world.timestep % self.subsidy_interval == 0:
+                masks[self.world.planner.idx] = self.default_planner_action_mask
+            else:
+                masks[self.world.planner.idx] = self.no_op_planner_action_mask
+        return masks
+
+    def get_data_dictionary(self):
+        """
+        Create a dictionary of data to push to the device
+        """
+        data_dict = DataFeed()
+        data_dict.add_data(
+            name="subsidy_interval",
+            data=self.subsidy_interval,
+        )
+        data_dict.add_data(
+            name="num_subsidy_levels",
+            data=self.num_subsidy_levels,
+        )
+        data_dict.add_data(
+            name="max_daily_subsidy_per_state",
+            data=self.max_daily_subsidy_per_state,
+        )
+        data_dict.add_data(
+            name="default_planner_action_mask",
+            data=[1] + self.default_planner_action_mask,
+        )
+        data_dict.add_data(
+            name="no_op_planner_action_mask",
+            data=[1] + self.no_op_planner_action_mask,
+        )
+        return data_dict
+
+    def get_tensor_dictionary(self):
+        """
+        Create a dictionary of (Pytorch-accessible) data to push to the device
+        """
+        tensor_dict = DataFeed()
+        return tensor_dict
+
+    def component_step(self):
+        if self.world.use_cuda:
+            self.world.cuda_component_step[self.name](
+                self.world.cuda_data_manager.device_data("subsidy_level"),
+                self.world.cuda_data_manager.device_data("subsidy"),
+                self.world.cuda_data_manager.device_data("subsidy_interval"),
+                self.world.cuda_data_manager.device_data("num_subsidy_levels"),
+                self.world.cuda_data_manager.device_data("max_daily_subsidy_per_state"),
+                self.world.cuda_data_manager.device_data("default_planner_action_mask"),
+                self.world.cuda_data_manager.device_data("no_op_planner_action_mask"),
+                self.world.cuda_data_manager.device_data(f"{_ACTIONS}_p"),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_a_{self.name}-t_until_next_subsidy"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_a_{self.name}-current_subsidy_level"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_p_{self.name}-t_until_next_subsidy"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_p_{self.name}-current_subsidy_level"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_p_action_mask"
+                ),
+                self.world.cuda_data_manager.device_data("_timestep_"),
+                self.world.cuda_data_manager.meta_info("n_agents"),
+                self.world.cuda_data_manager.meta_info("episode_length"),
+                block=self.world.cuda_function_manager.block,
+                grid=self.world.cuda_function_manager.grid,
+            )
+        else:
+            if self.world.use_real_world_policies:
+                if self._subsidy_amount_per_level is None:
+                    self._subsidy_amount_per_level = (
+                        self.world.us_population
+                        * self.max_annual_subsidy_per_person
+                        / self.num_subsidy_levels
+                        * self.subsidy_interval
+                        / 365
+                    )
+                    self._subsidy_level_array = np.zeros((self._episode_length + 1))
+                # Use the action taken in the previous timestep
+                current_subsidy_amount = self.world.real_world_subsidy[
+                    self.world.timestep - 1
+                ]
+                if current_subsidy_amount > 0:
+                    _subsidy_level = np.round(
+                        (current_subsidy_amount / self._subsidy_amount_per_level)
+                    )
+                    for t_idx in range(
+                        self.world.timestep - 1,
+                        min(
+                            len(self._subsidy_level_array),
+                            self.world.timestep - 1 + self.subsidy_interval,
+                        ),
+                    ):
+                        self._subsidy_level_array[t_idx] += _subsidy_level
+                subsidy_level = self._subsidy_level_array[self.world.timestep - 1]
+            else:
+                # Update the subsidy level only every self.subsidy_interval, since the
+                # other actions are masked out.
+                if (self.world.timestep - 1) % self.subsidy_interval == 0:
+                    subsidy_level = self.world.planner.get_component_action(self.name)
+                else:
+                    subsidy_level = self.world.planner.state["Current Subsidy Level"]
+
+            assert 0 <= subsidy_level <= self.num_subsidy_levels
+            self.world.planner.state["Current Subsidy Level"] = np.array(
+                subsidy_level
+            ).astype(self.np_int_dtype)
+
+            # Update subsidy level
+            subsidy_level_frac = subsidy_level / self.num_subsidy_levels
+            daily_statewise_subsidy = (
+                subsidy_level_frac * self.max_daily_subsidy_per_state
+            )
+
+            self.world.global_state["Subsidy"][
+                self.world.timestep
+            ] = daily_statewise_subsidy
+            self.world.planner.state["Total Subsidy"] += np.sum(daily_statewise_subsidy)
+
+    def generate_observations(self):
+        # Allow the agents/planner to know when the next subsidy might come.
+        # Obs should = 0 when the next timestep could include a subsidy
+        t_since_last_subsidy = self.world.timestep % self.subsidy_interval
+        # (this is normalized to 0<-->1)
+        t_until_next_subsidy = self.subsidy_interval - t_since_last_subsidy
+        t_vec = t_until_next_subsidy * np.ones(self.n_agents)
+
+        current_subsidy_level = self.world.planner.state["Current Subsidy Level"]
+        sl_vec = current_subsidy_level * np.ones(self.n_agents)
+
+        # Normalized observations
+        obs_dict = dict()
+        obs_dict["a"] = {
+            "t_until_next_subsidy": t_vec / self.subsidy_interval,
+            "current_subsidy_level": sl_vec / self.num_subsidy_levels,
+        }
+        obs_dict[self.world.planner.idx] = {
+            "t_until_next_subsidy": t_until_next_subsidy / self.subsidy_interval,
+            "current_subsidy_level": current_subsidy_level / self.num_subsidy_levels,
+        }
+
+        return obs_dict
+
+
+@component_registry.add
+class VaccinationCampaign(BaseComponent):
+    """
+    Implements a (passive) component for delivering vaccines to agents once a certain
+    amount of time has elapsed.
+
+    Args:
+        daily_vaccines_per_million_people (int): The number of vaccines available per
+            million people everyday.
+        delivery_interval (int): The number of days between vaccine deliveries.
+        vaccine_delivery_start_date (string): The date (YYYY-MM-DD) when the
+            vaccination begins.
+    """
+
+    name = "VaccinationCampaign"
+    required_entities = []
+    agent_subclasses = ["BasicMobileAgent"]
+
+    def __init__(
+        self,
+        *base_component_args,
+        daily_vaccines_per_million_people=4500,
+        delivery_interval=1,
+        vaccine_delivery_start_date="2020-12-22",
+        observe_rate=False,
+        **base_component_kwargs,
+    ):
+        self.daily_vaccines_per_million_people = int(daily_vaccines_per_million_people)
+        assert 0 <= self.daily_vaccines_per_million_people <= 1e6
+
+        self.delivery_interval = int(delivery_interval)
+        assert 1 <= self.delivery_interval <= 5000
+
+        try:
+            self.vaccine_delivery_start_date = datetime.strptime(
+                vaccine_delivery_start_date, "%Y-%m-%d"
+            )
+        except ValueError:
+            print("Incorrect data format, should be YYYY-MM-DD")
+
+        # (This will  be overwritten during component_step (see below))
+        self._time_when_vaccine_delivery_begins = None
+
+        self.np_int_dtype = np.int32
+
+        self.observe_rate = bool(observe_rate)
+
+        super().__init__(*base_component_args, **base_component_kwargs)
+
+        # (This will be overwritten during reset; see below)
+        self._num_vaccines_per_delivery = None
+        # Convenience for obs (see usage below):
+        self._t_first_delivery = None
+
+    @property
+    def num_vaccines_per_delivery(self):
+        if self._num_vaccines_per_delivery is None:
+            # Pre-compute dispersal numbers
+            millions_of_residents = self.world.us_state_population / 1e6
+            daily_vaccines = (
+                millions_of_residents * self.daily_vaccines_per_million_people
+            )
+            num_vaccines_per_delivery = np.floor(
+                self.delivery_interval * daily_vaccines
+            )
+            self._num_vaccines_per_delivery = np.array(
+                num_vaccines_per_delivery, dtype=self.np_int_dtype
+            )
+        return self._num_vaccines_per_delivery
+
+    @property
+    def time_when_vaccine_delivery_begins(self):
+        if self._time_when_vaccine_delivery_begins is None:
+            self._time_when_vaccine_delivery_begins = (
+                self.vaccine_delivery_start_date - self.world.start_date
+            ).days
+        return self._time_when_vaccine_delivery_begins
+
+    def get_additional_state_fields(self, agent_cls_name):
+        if agent_cls_name == "BasicMobileAgent":
+            return {"Total Vaccinated": 0, "Vaccines Available": 0}
+        return {}
+
+    def additional_reset_steps(self):
+        pass
+
+    def get_n_actions(self, agent_cls_name):
+        return  # Passive component
+
+    def generate_masks(self, completions=0):
+        return {}  # Passive component
+
+    def get_data_dictionary(self):
+        """
+        Create a dictionary of data to push to the device
+        """
+        data_dict = DataFeed()
+        data_dict.add_data(
+            name="num_vaccines_per_delivery",
+            data=self.num_vaccines_per_delivery,
+        )
+        data_dict.add_data(
+            name="delivery_interval",
+            data=self.delivery_interval,
+        )
+        data_dict.add_data(
+            name="time_when_vaccine_delivery_begins",
+            data=self.time_when_vaccine_delivery_begins,
+        )
+        data_dict.add_data(
+            name="num_vaccines_available_t",
+            data=np.zeros(self.n_agents),
+            save_copy_and_apply_at_reset=True,
+        )
+        return data_dict
+
+    def get_tensor_dictionary(self):
+        """
+        Create a dictionary of (Pytorch-accessible) data to push to the device
+        """
+        tensor_dict = DataFeed()
+        return tensor_dict
+
+    def component_step(self):
+        if self.world.use_cuda:
+            self.world.cuda_component_step[self.name](
+                self.world.cuda_data_manager.device_data("vaccinated"),
+                self.world.cuda_data_manager.device_data("num_vaccines_per_delivery"),
+                self.world.cuda_data_manager.device_data("num_vaccines_available_t"),
+                self.world.cuda_data_manager.device_data("delivery_interval"),
+                self.world.cuda_data_manager.device_data(
+                    "time_when_vaccine_delivery_begins"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_a_{self.name}-t_until_next_vaccines"
+                ),
+                self.world.cuda_data_manager.device_data(
+                    f"{_OBSERVATIONS}_p_{self.name}-t_until_next_vaccines"
+                ),
+                self.world.cuda_data_manager.device_data("_timestep_"),
+                self.world.cuda_data_manager.meta_info("n_agents"),
+                self.world.cuda_data_manager.meta_info("episode_length"),
+                block=self.world.cuda_function_manager.block,
+                grid=self.world.cuda_function_manager.grid,
+            )
+        else:
+            # Do nothing if vaccines are not available yet
+            if self.world.timestep < self.time_when_vaccine_delivery_begins:
+                return
+
+            # Do nothing if this is not the start of a delivery interval.
+            # Vaccines are delivered at the start of each interval.
+            if (self.world.timestep % self.delivery_interval) != 0:
+                return
+
+            # Deliver vaccines to each state
+            for aidx, vaccines in enumerate(self.num_vaccines_per_delivery):
+                self.world.agents[aidx].state["Vaccines Available"] += vaccines
+
+    def generate_observations(self):
+        # Allow the agents/planner to know when the next vaccines might come.
+        # Obs should = 0 when the next timestep will deliver vaccines
+        # (this is normalized to 0<-->1)
+
+        if self._t_first_delivery is None:
+            self._t_first_delivery = int(self.time_when_vaccine_delivery_begins)
+            while (self._t_first_delivery % self.delivery_interval) != 0:
+                self._t_first_delivery += 1
+
+        next_t = self.world.timestep + 1
+        if next_t <= self._t_first_delivery:
+            t_until_next_vac = np.minimum(
+                1, (self._t_first_delivery - next_t) / self.delivery_interval
+            )
+            next_vax_rate = 0.0
+        else:
+            t_since_last_vac = next_t % self.delivery_interval
+            t_until_next_vac = self.delivery_interval - t_since_last_vac
+            next_vax_rate = self.daily_vaccines_per_million_people / 1e6
+        t_vec = t_until_next_vac * np.ones(self.n_agents)
+        r_vec = next_vax_rate * np.ones(self.n_agents)
+
+        # Normalized observations
+        obs_dict = dict()
+        obs_dict["a"] = {"t_until_next_vaccines": t_vec / self.delivery_interval}
+        obs_dict[self.world.planner.idx] = {
+            "t_until_next_vaccines": t_until_next_vac / self.delivery_interval
+        }
+
+        if self.observe_rate:
+            obs_dict["a"]["next_vaccination_rate"] = r_vec
+            obs_dict["p"]["next_vaccination_rate"] = float(next_vax_rate)
+
+        return obs_dict

ai_economist/foundation/components/covid19_components_step.cu

@@ -0,0 +1,263 @@
+// Copyright (c) 2021, salesforce.com, inc.
+// All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+// For full license text, see the LICENSE file in the repo root
+// or https://opensource.org/licenses/BSD-3-Clause
+
+extern "C" {
+    // CUDA version of the components in
+    // "ai_economist.foundation.components.covid19_components.py"
+    __global__ void CudaControlUSStateOpenCloseStatusStep(
+        int * stringency_level,
+        const int kActionCooldownPeriod,
+        int * action_in_cooldown_until,
+        const int * kDefaultAgentActionMask,
+        const int * kNoOpAgentActionMask,
+        const int kNumStringencyLevels,
+        int * actions,
+        float * obs_a_stringency_policy_indicators,
+        float * obs_a_action_mask,
+        float * obs_p_stringency_policy_indicators,
+        int * env_timestep_arr,
+        const int kNumAgents,
+        const int kEpisodeLength
+    ) {
+        const int kEnvId = blockIdx.x;
+        const int kAgentId = threadIdx.x;
+
+        // Increment time ONCE -- only 1 thread can do this.
+        if (kAgentId == 0) {
+            env_timestep_arr[kEnvId] += 1;
+        }
+
+        // Wait here until timestep has been updated
+        __syncthreads();
+
+        assert(env_timestep_arr[kEnvId] > 0 &&
+            env_timestep_arr[kEnvId] <= kEpisodeLength);
+        assert (kAgentId <= kNumAgents - 1);
+
+        // Update the stringency levels for the US states
+        if (kAgentId < (kNumAgents - 1)) {
+            // Indices for time-dependent and time-independent arrays
+            // Time dependent arrays have shapes
+            // (num_envs, kEpisodeLength + 1, kNumAgents - 1)
+            // Time independent arrays have shapes (num_envs, kNumAgents - 1)
+            const int kArrayIdxOffset = kEnvId * (kEpisodeLength + 1) *
+                (kNumAgents - 1);
+            int time_dependent_array_index_curr_t = kArrayIdxOffset +
+                env_timestep_arr[kEnvId] * (kNumAgents - 1) + kAgentId;
+            int time_dependent_array_index_prev_t = kArrayIdxOffset +
+                (env_timestep_arr[kEnvId] - 1) * (kNumAgents - 1) + kAgentId;
+            const int time_independent_array_index = kEnvId * (kNumAgents - 1) +
+                kAgentId;
+
+            // action is not a NO-OP
+            if (actions[time_independent_array_index] != 0) {
+                stringency_level[time_dependent_array_index_curr_t] =
+                    actions[time_independent_array_index];
+            } else {
+                stringency_level[time_dependent_array_index_curr_t] =
+                    stringency_level[time_dependent_array_index_prev_t];
+            }
+
+            if (env_timestep_arr[kEnvId] == action_in_cooldown_until[
+                time_independent_array_index] + 1) {
+                if (actions[time_independent_array_index] != 0) {
+                    assert(0 <= actions[time_independent_array_index] <=
+                        kNumStringencyLevels);
+                    action_in_cooldown_until[time_independent_array_index] +=
+                        kActionCooldownPeriod;
+                } else {
+                    action_in_cooldown_until[time_independent_array_index] += 1;
+                }
+            }
+
+            obs_a_stringency_policy_indicators[
+                time_independent_array_index
+            ] = stringency_level[time_dependent_array_index_curr_t] /
+                static_cast<float>(kNumStringencyLevels);
+
+            // CUDA version of generate_masks()
+            for (int action_id = 0; action_id < (kNumStringencyLevels + 1);
+                action_id++) {
+                int action_mask_array_index =
+                    kEnvId * (kNumStringencyLevels + 1) *
+                    (kNumAgents - 1) + action_id * (kNumAgents - 1) + kAgentId;
+                if (env_timestep_arr[kEnvId] < action_in_cooldown_until[
+                    time_independent_array_index]
+                ) {
+                    obs_a_action_mask[action_mask_array_index] =
+                    kNoOpAgentActionMask[action_id];
+                } else {
+                    obs_a_action_mask[action_mask_array_index] =
+                    kDefaultAgentActionMask[action_id];
+                }
+            }
+        }
+
+        // Update planner obs after all the agents' obs are updated
+        __syncthreads();
+
+        if (kAgentId == kNumAgents - 1) {
+            for (int ag_id = 0; ag_id < (kNumAgents - 1); ag_id++) {
+                const int kIndex = kEnvId * (kNumAgents - 1) + ag_id;
+                obs_p_stringency_policy_indicators[
+                    kIndex
+                ] = 
+                    obs_a_stringency_policy_indicators[
+                        kIndex
+                    ];
+            }
+        }
+    }
+
+    __global__ void CudaFederalGovernmentSubsidyStep(
+        int * subsidy_level,
+        float * subsidy,
+        const int kSubsidyInterval,
+        const int kNumSubsidyLevels,
+        const float * KMaxDailySubsidyPerState,
+        const int * kDefaultPlannerActionMask,
+        const int * kNoOpPlannerActionMask,
+        int * actions,
+        float * obs_a_time_until_next_subsidy,
+        float * obs_a_current_subsidy_level,
+        float * obs_p_time_until_next_subsidy,
+        float * obs_p_current_subsidy_level,
+        float * obs_p_action_mask,
+        int * env_timestep_arr,
+        const int kNumAgents,
+        const int kEpisodeLength
+    ) {
+        const int kEnvId = blockIdx.x;
+        const int kAgentId = threadIdx.x;
+
+        assert(env_timestep_arr[kEnvId] > 0 &&
+            env_timestep_arr[kEnvId] <= kEpisodeLength);
+        assert (kAgentId <= kNumAgents - 1);
+
+        int t_since_last_subsidy = env_timestep_arr[kEnvId] %
+            kSubsidyInterval;
+
+        // Setting the (federal government) planner's subsidy level
+        // to be the subsidy level for all the US states
+        if (kAgentId < kNumAgents - 1) {
+            // Indices for time-dependent and time-independent arrays
+            // Time dependent arrays have shapes (num_envs,
+            // kEpisodeLength + 1, kNumAgents - 1)
+            // Time independent arrays have shapes (num_envs, kNumAgents - 1)
+            const int kArrayIdxOffset = kEnvId * (kEpisodeLength + 1) *
+                (kNumAgents - 1);
+            int time_dependent_array_index_curr_t = kArrayIdxOffset +
+                env_timestep_arr[kEnvId] * (kNumAgents - 1) + kAgentId;
+            int time_dependent_array_index_prev_t = kArrayIdxOffset +
+                (env_timestep_arr[kEnvId] - 1) * (kNumAgents - 1) + kAgentId;
+            const int time_independent_array_index = kEnvId *
+                (kNumAgents - 1) + kAgentId;
+
+            if ((env_timestep_arr[kEnvId] - 1) % kSubsidyInterval == 0) {
+                assert(0 <= actions[kEnvId] <= kNumSubsidyLevels);
+                subsidy_level[time_dependent_array_index_curr_t] =
+                    actions[kEnvId];
+            } else {
+                subsidy_level[time_dependent_array_index_curr_t] =
+                    subsidy_level[time_dependent_array_index_prev_t];
+            }
+            // Setting the subsidies for the US states
+            // based on the federal government's subsidy level
+            subsidy[time_dependent_array_index_curr_t] =
+                subsidy_level[time_dependent_array_index_curr_t] *
+                KMaxDailySubsidyPerState[kAgentId] / kNumSubsidyLevels;
+
+            obs_a_time_until_next_subsidy[
+                time_independent_array_index] =
+                    1 - (t_since_last_subsidy /
+                    static_cast<float>(kSubsidyInterval));
+            obs_a_current_subsidy_level[
+                time_independent_array_index] =
+                    subsidy_level[time_dependent_array_index_curr_t] /
+                    static_cast<float>(kNumSubsidyLevels);
+        } else if (kAgentId == (kNumAgents - 1)) {
+            for (int action_id = 0; action_id < kNumSubsidyLevels + 1;
+                action_id++) {
+                int action_mask_array_index = kEnvId *
+                    (kNumSubsidyLevels + 1) + action_id;
+                if (env_timestep_arr[kEnvId] % kSubsidyInterval == 0) {
+                    obs_p_action_mask[action_mask_array_index] =
+                        kDefaultPlannerActionMask[action_id];
+                } else {
+                    obs_p_action_mask[action_mask_array_index] =
+                        kNoOpPlannerActionMask[action_id];
+                }
+            }
+            // Update planner obs after the agent's obs are updated
+            __syncthreads();
+
+            if (kAgentId == (kNumAgents - 1)) {
+                // Just use the values for agent id 0
+                obs_p_time_until_next_subsidy[kEnvId] =
+                    obs_a_time_until_next_subsidy[
+                        kEnvId * (kNumAgents - 1)
+                    ];
+                obs_p_current_subsidy_level[kEnvId] = 
+                    obs_a_current_subsidy_level[
+                        kEnvId * (kNumAgents - 1)
+                    ];
+            }
+        }
+    }
+
+    __global__ void CudaVaccinationCampaignStep(
+        int * vaccinated,
+        const int * kNumVaccinesPerDelivery,
+        int * num_vaccines_available_t,
+        const int kDeliveryInterval,
+        const int kTimeWhenVaccineDeliveryBegins,
+        float * obs_a_vaccination_campaign_t_until_next_vaccines,
+        float * obs_p_vaccination_campaign_t_until_next_vaccines,
+        int * env_timestep_arr,
+        int kNumAgents,
+        int kEpisodeLength
+    ) {
+        const int kEnvId = blockIdx.x;
+        const int kAgentId = threadIdx.x;
+
+        assert(env_timestep_arr[kEnvId] > 0 && env_timestep_arr[kEnvId] <=
+            kEpisodeLength);
+        assert(kTimeWhenVaccineDeliveryBegins > 0);
+        assert (kAgentId <= kNumAgents - 1);
+
+        // CUDA version of generate observations()
+        int t_first_delivery = kTimeWhenVaccineDeliveryBegins +
+            kTimeWhenVaccineDeliveryBegins % kDeliveryInterval;
+        int next_t = env_timestep_arr[kEnvId] + 1;
+        float t_until_next_vac;
+        if (next_t <= t_first_delivery) {
+            t_until_next_vac = min(
+                1,
+                (t_first_delivery - next_t) / kDeliveryInterval);
+        } else {
+            float t_since_last_vac = next_t % kDeliveryInterval;
+            t_until_next_vac = 1 - (t_since_last_vac / kDeliveryInterval);
+        }
+
+        // Update the vaccinated numbers for just the US states
+        if (kAgentId < (kNumAgents - 1)) {
+            const int time_independent_array_index = kEnvId *
+                (kNumAgents - 1) + kAgentId;
+            if ((env_timestep_arr[kEnvId] >= kTimeWhenVaccineDeliveryBegins) &&
+                (env_timestep_arr[kEnvId] % kDeliveryInterval == 0)) {
+                num_vaccines_available_t[time_independent_array_index] =
+                    kNumVaccinesPerDelivery[kAgentId];
+            } else {
+                num_vaccines_available_t[time_independent_array_index] = 0;
+            }
+            obs_a_vaccination_campaign_t_until_next_vaccines[
+                time_independent_array_index] = t_until_next_vac;
+        } else if (kAgentId == kNumAgents - 1) {
+            obs_p_vaccination_campaign_t_until_next_vaccines[kEnvId] =
+            t_until_next_vac;
+        }
+    }
+}

ai_economist/foundation/components/move.py

@@ -0,0 +1,222 @@
+# Copyright (c) 2020, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root
+# or https://opensource.org/licenses/BSD-3-Clause
+
+import numpy as np
+from numpy.random import rand
+
+from ai_economist.foundation.base.base_component import (
+    BaseComponent,
+    component_registry,
+)
+
+
+@component_registry.add
+class Gather(BaseComponent):
+    """
+    Allows mobile agents to move around the world and collect resources and prevents
+    agents from moving to invalid locations.
+
+    Can be configured to include collection skill, where agents have heterogeneous
+    probabilities of collecting bonus resources without additional labor cost.
+
+    Args:
+        move_labor (float): Labor cost associated with movement. Must be >= 0.
+            Default is 1.0.
+        collect_labor (float): Labor cost associated with collecting resources. This
+            cost is added (in addition to any movement cost) when the agent lands on
+            a tile that is populated with resources (triggering collection).
+            Must be >= 0. Default is 1.0.
+        skill_dist (str): Distribution type for sampling skills. Default ("none")
+            gives all agents identical skill equal to a bonus prob of 0. "pareto" and
+            "lognormal" sample skills from the associated distributions.
+    """
+
+    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)
+        assert self.move_labor >= 0
+
+        self.collect_labor = float(collect_labor)
+        assert self.collect_labor >= 0
+
+        self.skill_dist = skill_dist.lower()
+        assert self.skill_dist in ["none", "pareto", "lognormal"]
+
+        self.gathers = []
+
+        self._aidx = np.arange(self.n_agents)[:, None].repeat(4, axis=1)
+        self._roff = np.array([[0, 0, -1, 1]])
+        self._coff = np.array([[-1, 1, 0, 0]])
+
+    # Required methods for implementing components
+    # --------------------------------------------
+
+    def get_n_actions(self, agent_cls_name):
+        """
+        See base_component.py for detailed description.
+
+        Adds 4 actions (move up, down, left, or right) for mobile agents.
+        """
+        # This component adds 4 action that agents can take:
+        # move up, down, left, or right
+        if agent_cls_name == "BasicMobileAgent":
+            return 4
+        return None
+
+    def get_additional_state_fields(self, agent_cls_name):
+        """
+        See base_component.py for detailed description.
+
+        For mobile agents, add state field for collection skill.
+        """
+        if agent_cls_name not in self.agent_subclasses:
+            return {}
+        if agent_cls_name == "BasicMobileAgent":
+            return {"bonus_gather_prob": 0.0}
+        raise NotImplementedError
+
+    def component_step(self):
+        """
+        See base_component.py for detailed description.
+
+        Move to adjacent, unoccupied locations. Collect resources when moving to
+        populated resource tiles, adding the resource to the agent's inventory and
+        de-populating it from the tile.
+        """
+        world = self.world
+
+        gathers = []
+        for agent in world.get_random_order_agents():
+
+            if self.name not in agent.action:
+                return
+            action = agent.get_component_action(self.name)
+
+            r, c = [int(x) for x in agent.loc]
+
+            if action == 0:  # NO-OP!
+                new_r, new_c = r, c
+
+            elif action <= 4:
+                if action == 1:  # Left
+                    new_r, new_c = r, c - 1
+                elif action == 2:  # Right
+                    new_r, new_c = r, c + 1
+                elif action == 3:  # Up
+                    new_r, new_c = r - 1, c
+                else:  # action == 4, # Down
+                    new_r, new_c = r + 1, c
+
+                # Attempt to move the agent (if the new coordinates aren't accessible,
+                # nothing will happen)
+                new_r, new_c = world.set_agent_loc(agent, new_r, new_c)
+
+                # If the agent did move, incur the labor cost of moving
+                if (new_r != r) or (new_c != c):
+                    agent.state["endogenous"]["Labor"] += self.move_labor
+
+            else:
+                raise ValueError
+
+            for resource, health in world.location_resources(new_r, new_c).items():
+                if health >= 1:
+                    n_gathered = 1 + (rand() < agent.state["bonus_gather_prob"])
+                    agent.state["inventory"][resource] += n_gathered
+                    world.consume_resource(resource, new_r, new_c)
+                    # Incur the labor cost of collecting a resource
+                    agent.state["endogenous"]["Labor"] += self.collect_labor
+                    # Log the gather
+                    gathers.append(
+                        dict(
+                            agent=agent.idx,
+                            resource=resource,
+                            n=n_gathered,
+                            loc=[new_r, new_c],
+                        )
+                    )
+
+        self.gathers.append(gathers)
+
+    def generate_observations(self):
+        """
+        See base_component.py for detailed description.
+
+        Here, agents observe their collection skill. The planner does not observe
+        anything from this component.
+        """
+        return {
+            str(agent.idx): {"bonus_gather_prob": agent.state["bonus_gather_prob"]}
+            for agent in self.world.agents
+        }
+
+    def generate_masks(self, completions=0):
+        """
+        See base_component.py for detailed description.
+
+        Prevent moving to adjacent tiles that are already occupied (or outside the
+        boundaries of the world)
+        """
+        world = self.world
+
+        coords = np.array([agent.loc for agent in world.agents])[:, :, None]
+        ris = coords[:, 0] + self._roff + 1
+        cis = coords[:, 1] + self._coff + 1
+
+        occ = np.pad(world.maps.unoccupied, ((1, 1), (1, 1)))
+        acc = np.pad(world.maps.accessibility, ((0, 0), (1, 1), (1, 1)))
+        mask_array = np.logical_and(occ[ris, cis], acc[self._aidx, ris, cis]).astype(
+            np.float32
+        )
+
+        masks = {agent.idx: mask_array[i] for i, agent in enumerate(world.agents)}
+
+        return masks
+
+    # For non-required customization
+    # ------------------------------
+
+    def additional_reset_steps(self):
+        """
+        See base_component.py for detailed description.
+
+        Re-sample agents' collection skills.
+        """
+        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
+            else:
+                raise NotImplementedError
+            agent.state["bonus_gather_prob"] = float(bonus_rate)
+
+        self.gathers = []
+
+    def get_dense_log(self):
+        """
+        Log resource collections.
+
+        Returns:
+            gathers (list): A list of gather events. Each entry corresponds to a single
+                timestep and contains a description of any resource gathers that
+                occurred on that timestep.
+
+        """
+        return self.gathers

ai_economist/foundation/components/simple_labor.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2021, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root
+# or https://opensource.org/licenses/BSD-3-Clause
+
+import numpy as np
+
+from ai_economist.foundation.base.base_component import (
+    BaseComponent,
+    component_registry,
+)
+
+
+@component_registry.add
+class SimpleLabor(BaseComponent):
+    """
+    Allows Agents to select a level of labor, which earns income based on skill.
+
+    Labor is "simple" because this simplifies labor to a choice along a 1D axis. More
+    concretely, this component adds 100 labor actions, each representing a choice of
+    how many hours to work, e.g. action 50 represents doing 50 hours of work; each
+    Agent earns income proportional to the product of its labor amount (representing
+    hours worked) and its skill (representing wage), with higher skill and higher labor
+    yielding higher income.
+
+    This component is intended to be used with the 'PeriodicBracketTax' component and
+    the 'one-step-economy' scenario.
+
+    Args:
+        mask_first_step (bool): Defaults to True. If True, masks all non-0 labor
+            actions on the first step of the environment. When combined with the
+            intended component/scenario, the first env step is used to set taxes
+            (via the 'redistribution' component) and the second step is used to
+            select labor (via this component).
+        payment_max_skill_multiplier (float): When determining the skill level of
+            each Agent, sampled skills are clipped to this maximum value.
+    """
+
+    name = "SimpleLabor"
+    required_entities = ["Coin"]
+    agent_subclasses = ["BasicMobileAgent"]
+
+    def __init__(
+        self,
+        *base_component_args,
+        mask_first_step=True,
+        payment_max_skill_multiplier=3,
+        pareto_param=4.0,
+        **base_component_kwargs
+    ):
+        super().__init__(*base_component_args, **base_component_kwargs)
+
+        # This defines the size of the action space (the max # hours an agent can work).
+        self.num_labor_hours = 100  # max 100 hours
+
+        assert isinstance(mask_first_step, bool)
+        self.mask_first_step = mask_first_step
+
+        self.is_first_step = True
+        self.common_mask_on = {
+            agent.idx: np.ones((self.num_labor_hours,)) for agent in self.world.agents
+        }
+        self.common_mask_off = {
+            agent.idx: np.zeros((self.num_labor_hours,)) for agent in self.world.agents
+        }
+
+        # Skill distribution
+        self.pareto_param = float(pareto_param)
+        assert self.pareto_param > 0
+        self.payment_max_skill_multiplier = float(payment_max_skill_multiplier)
+        pmsm = self.payment_max_skill_multiplier
+        num_agents = len(self.world.agents)
+        # Generate a batch (1000) of num_agents (sorted/clipped) Pareto samples.
+        pareto_samples = np.random.pareto(4, size=(1000, num_agents))
+        clipped_skills = np.minimum(pmsm, (pmsm - 1) * pareto_samples + 1)
+        sorted_clipped_skills = np.sort(clipped_skills, axis=1)
+        # The skill level of the i-th skill-ranked agent is the average of the
+        # i-th ranked samples throughout the batch.
+        self.skills = sorted_clipped_skills.mean(axis=0)
+
+    def get_additional_state_fields(self, agent_cls_name):
+        if agent_cls_name == "BasicMobileAgent":
+            return {"skill": 0, "production": 0}
+        return {}
+
+    def additional_reset_steps(self):
+        self.is_first_step = True
+        for agent in self.world.agents:
+            agent.state["skill"] = self.skills[agent.idx]
+
+    def get_n_actions(self, agent_cls_name):
+        if agent_cls_name == "BasicMobileAgent":
+            return self.num_labor_hours
+        return None
+
+    def generate_masks(self, completions=0):
+        if self.is_first_step:
+            self.is_first_step = False
+            if self.mask_first_step:
+                return self.common_mask_off
+
+        return self.common_mask_on
+
+    def component_step(self):
+
+        for agent in self.world.get_random_order_agents():
+
+            action = agent.get_component_action(self.name)
+
+            if action == 0:  # NO-OP.
+                # Agent is not interacting with this component.
+                continue
+
+            if 1 <= action <= self.num_labor_hours:  # set reopening phase
+
+                hours_worked = action  # NO-OP is 0 hours.
+                agent.state["endogenous"]["Labor"] = hours_worked
+
+                payoff = hours_worked * agent.state["skill"]
+                agent.state["production"] += payoff
+                agent.inventory["Coin"] += payoff
+
+            else:
+                # If action > num_labor_hours, this is an error.
+                raise ValueError
+
+    def generate_observations(self):
+        obs_dict = dict()
+        for agent in self.world.agents:
+            obs_dict[str(agent.idx)] = {
+                "skill": agent.state["skill"] / self.payment_max_skill_multiplier
+            }
+        return obs_dict

ai_economist/foundation/components/utils.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2020, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root
+# or https://opensource.org/licenses/BSD-3-Clause
+
+import numpy as np
+
+
+def annealed_tax_limit(completions, warmup_period, slope, final_max_tax_value=1.0):
+    """
+    Compute the maximum tax rate available at this stage of tax annealing.
+
+    This function uses the number of episode completions and the annealing schedule
+    (warmup_period, slope, & final_max_tax_value) to determine what the maximum tax
+    rate can be.
+    This type of annealing allows for a tax curriculum where earlier episodes are
+    restricted to lower tax rates. As more episodes are played, higher tax values are
+    allowed.
+
+    Args:
+        completions (int): Number of times the environment has completed an episode.
+            Expected to be >= 0.
+        warmup_period (int): Until warmup_period completions, only allow 0 tax. Using
+            a negative value will enable non-0 taxes at 0 environment completions.
+        slope (float): After warmup_period completions, percentage of full tax value
+            unmasked with each new completion.
+        final_max_tax_value (float): The maximum tax value at the end of annealing.
+
+    Returns:
+        A scalar value indicating the maximum tax at this stage of annealing.
+
+    Example:
+        >> WARMUP = 100
+        >> SLOPE = 0.01
+        >> annealed_tax_limit(0, WARMUP, SLOPE)
+        0.0
+        >> annealed_tax_limit(100, WARMUP, SLOPE)
+        0.0
+        >> annealed_tax_limit(150, WARMUP, SLOPE)
+        0.5
+        >> annealed_tax_limit(200, WARMUP, SLOPE)
+        1.0
+        >> annealed_tax_limit(1000, WARMUP, SLOPE)
+        1.0
+    """
+    # What percentage of the full range is currently visible
+    # (between 0 [only 0 tax] and 1 [all taxes visible])
+    percentage_visible = np.maximum(
+        0.0, np.minimum(1.0, slope * (completions - warmup_period))
+    )
+
+    # Determine the highest allowable tax,
+    # given the current position in the annealing schedule
+    current_max_tax = percentage_visible * final_max_tax_value
+
+    return current_max_tax
+
+
+def annealed_tax_mask(completions, warmup_period, slope, tax_values):
+    """
+    Generate a mask applied to a set of tax values for the purpose of tax annealing.
+
+    This function uses the number of episode completions and the annealing schedule
+    to determine which of the tax values are considered valid. The most extreme
+    tax/subsidy values are unmasked last. Zero tax is always unmasked (i.e. always
+    valid).
+    This type of annealing allows for a tax curriculum where earlier episodes are
+    restricted to lower tax rates. As more episodes are played, higher tax values are
+    allowed.
+
+    Args:
+        completions (int): Number of times the environment has completed an episode.
+            Expected to be >= 0.
+        warmup_period (int): Until warmup_period completions, only allow 0 tax. Using
+            a negative value will enable non-0 taxes at 0 environment completions.
+        slope (float): After warmup_period completions, percentage of full tax value
+            unmasked with each new completion.
+        tax_values (list): The list of tax values associated with each action to
+            which this mask will apply.
+
+    Returns:
+        A binary mask with same shape as tax_values, indicating which tax values are
+            currently valid.
+
+    Example:
+        >> WARMUP = 100
+        >> SLOPE = 0.01
+        >> TAX_VALUES = [0.0, 0.25, 0.50, 0.75, 1.0]
+        >> annealed_tax_limit(0, WARMUP, SLOPE, TAX_VALUES)
+        [0, 0, 0, 0, 0]
+        >> annealed_tax_limit(100, WARMUP, SLOPE, TAX_VALUES)
+        [0, 0, 0, 0, 0]
+        >> annealed_tax_limit(150, WARMUP, SLOPE, TAX_VALUES)
+        [1, 1, 1, 0, 0]
+        >> annealed_tax_limit(200, WARMUP, SLOPE, TAX_VALUES)
+        [1, 1, 1, 1, 1]
+        >> annealed_tax_limit(1000, WARMUP, SLOPE, TAX_VALUES)
+        [1, 1, 1, 1, 1]
+    """
+    # Infer the most extreme tax level from the supplied tax values.
+    abs_tax = np.abs(tax_values)
+    full_tax_amount = np.max(abs_tax)
+
+    # Determine the highest allowable tax, given the current position
+    # in the annealing schedule
+    max_absolute_visible_tax = annealed_tax_limit(
+        completions, warmup_period, slope, full_tax_amount
+    )
+
+    # Return a binary mask to allow for taxes
+    # at or below the highest absolute visible tax
+    return np.less_equal(np.abs(tax_values), max_absolute_visible_tax).astype(
+        np.float32
+    )