a not-too-general value backwards induction algorithm built on DBlocks

Documentation/CHANGELOG.md

@@ -8,6 +8,17 @@ For more information on HARK, see [our Github organization](https://github.com/e
 
 ## Changes
 
+### 0.16.0
+
+Under development
+
+#### Major Changes
+
+- Adds `HARK.algos.vbi` as a general algorithm for solving the optimization step of an agent's problem. [#1438](https://github.com/econ-ark/HARK/pull/1438)
+
+#### Minor Changes
+
+
 ### 0.15.0
 
 Release Date: June 4, 2024
@@ -16,7 +27,7 @@ Note: Due to major changes on this release, you may need to adjust how AgentType
 
 This release drops support for Python 3.8 and 3.9, consistent with SPEC 0, and adds support for Python 3.11 and 3.12. We expect that all HARK features still work with the older versions, but they are no longer part of our testing regimen.
 
-### Major Changes
+#### Major Changes
 
 - Drop official support for Python 3.8 and 3.9, add support for 3.11 and 3.12. [#1415](https://github.com/econ-ark/HARK/pull/1415)
 - Replace object-oriented solvers with single function versions. [#1394](https://github.com/econ-ark/HARK/pull/1394)
@@ -28,7 +39,7 @@ This release drops support for Python 3.8 and 3.9, consistent with SPEC 0, and a
 - Such constructed inputs can use alternate parameterizations / formats by changing the `constructor` function and providing its arguments in `parameters`.
 - Move `HARK.datasets` to `HARK.Calibration` for better organization of data and calibration tools. [#1430](https://github.com/econ-ark/HARK/pull/1430)
 
-### Minor Changes
+#### Minor Changes
 
 - Add option to pass pre-built grid to `LinearFast`. [#1388](https://github.com/econ-ark/HARK/pull/1388)
 - Moves calculation of stable points out of ConsIndShock solver, into method called by post_solve [#1349](https://github.com/econ-ark/HARK/pull/1349)

Documentation/reference/index.rst

@@ -5,6 +5,7 @@ API Reference
    :caption: Tools
    :maxdepth: 1
 
+   tools/algos
    tools/core
    tools/dcegm
    tools/distribution

Documentation/reference/tools/algos.rst

@@ -0,0 +1,7 @@
+algos
+--------
+
+.. toctree::
+   :maxdepth: 3
+
+   algos/vbi

Documentation/reference/tools/algos/vbi.rst

@@ -0,0 +1,7 @@
+algos.vbi
+----------
+
+.. automodule:: HARK.algos.vbi
+   :members:
+   :undoc-members:
+   :show-inheritance:

Documentation/reference/tools/index.rst.orig

@@ -0,0 +1,27 @@
+:orphan:
+
+Tools
+=====
+
+<<<<<<< HEAD
+.. toctree::
+   :maxdepth: 3
+
+   algos
+   core
+   dcegm
+   distribution
+   econforgeinterp
+   estimation
+   frame
+   helpers
+   interpolation
+   numba
+   parallel
+   rewards
+   simulation
+   utilities
+   validators
+=======
+See :doc:`../index`.
+>>>>>>> master

HARK/algos/tests/test_vbi.py

@@ -0,0 +1,34 @@
+import HARK.algos.vbi as vbi
+from HARK.distribution import Bernoulli
+from HARK.model import Control, DBlock
+import numpy as np
+
+import unittest
+
+
+block_1 = DBlock(
+    **{
+        "name": "vbi_test_1",
+        "shocks": {
+            "coin": Bernoulli(p=0.5),
+        },
+        "dynamics": {
+            "m": lambda y, coin: y + coin,
+            "c": Control(["m"]),
+            "a": lambda m, c: m - c,
+        },
+        "reward": {"u": lambda c: 1 - (c - 1) ** 2},
+    }
+)
+
+
+class test_vbi(unittest.TestCase):
+    # def setUp(self):
+    #    pass
+
+    def test_solve_block_1(self):
+        state_grid = {"m": np.linspace(0, 2, 10)}
+
+        dr, dec_vf, arr_vf = vbi.vbi_solve(block_1, lambda a: a, state_grid)
+
+        self.assertAlmostEqual(dr["c"](**{"m": 1}), 0.5)

HARK/algos/vbi.py

@@ -0,0 +1,162 @@
+"""
+Use backwards induction to derive the arrival value function
+from a continuation value function and stage dynamics.
+"""
+
+from HARK.model import DBlock
+import itertools
+import numpy as np
+from scipy.optimize import minimize
+from typing import Mapping, Sequence
+import xarray as xr
+
+
+def get_action_rule(action):
+    """
+    Produce a function from any inputs to a given value.
+    This is useful for constructing decision rules with fixed actions.
+    """
+
+    def ar():
+        return action
+
+    return ar
+
+
+def ar_from_data(da):
+    """
+    Produce a function from any inputs to a given value.
+    This is useful for constructing decision rules with fixed actions.
+    """
+
+    def ar(**args):
+        return da.interp(**args).values.tolist()
+
+    return ar
+
+
+Grid = Mapping[str, Sequence]
+
+
+def grid_to_data_array(
+    grid: Grid = {},  ## TODO: Better data structure here.
+):
+    """
+    Construct a zero-valued DataArray with the coordinates
+    based on the Grid passed in.
+
+    Parameters
+    ----------
+    grid: Grid
+        A mapping from variable labels to a sequence of numerical values.
+
+    Returns
+    --------
+    da xarray.DataArray
+        An xarray.DataArray with coordinates given by both grids.
+    """
+
+    coords = {**grid}
+
+    da = xr.DataArray(
+        np.empty([len(v) for v in coords.values()]), dims=coords.keys(), coords=coords
+    )
+
+    return da
+
+
+def vbi_solve(
+    block: DBlock, continuation, state_grid: Grid, disc_params={}, calibration={}
+):
+    """
+    Solve a DBlock using backwards induction on the value function.
+
+    Parameters
+    -----------
+    block
+    continuation
+
+    state_grid: Grid
+        This is a grid over all variables that the optimization will range over.
+        This should be just the information set of the decision variables.
+
+    disc_params
+    calibration
+    """
+
+    # state-rule value function
+    srv_function = block.get_state_rule_value_function_from_continuation(continuation)
+
+    controls = block.get_controls()
+
+    # pseudo
+    policy_data = grid_to_data_array(state_grid)
+    value_data = grid_to_data_array(state_grid)
+
+    # loop through every point in the state grid
+    for state_point in itertools.product(*state_grid.values()):
+        # build a dictionary from these states, as scope for the optimization
+        state_vals = {k: v for k, v in zip(state_grid.keys(), state_point)}
+
+        # The value of the action is computed given
+        # the problem calibration and the states for the current point on the
+        # state-grid.
+        pre_states = calibration.copy()
+        pre_states.update(state_vals)
+
+        # prepare function to optimize
+        def negated_value(a):  # old! (should be negative)
+            dr = {c: get_action_rule(a[i]) for i, c in enumerate(controls)}
+
+            # negative, for minimization later
+            return -srv_function(pre_states, dr)
+
+        ## get lower bound.
+        ## not yet implemented
+        lower_bound = np.array([-1e-12] * len(controls))  ## a really low number!
+
+        ## get upper bound
+        ## not yet implemented
+        upper_bound = np.array([1e12] * len(controls))
+
+        # pseudo
+        # optimize_action(pre_states, srv_function)
+
+        res = minimize(  # choice of
+            negated_value,
+            0,  # x0 is starting guess, here arbitrary.
+        )
+        print(res)
+
+        dr_best = {c: get_action_rule(res.x[i]) for i, c in enumerate(controls)}
+
+        if res.success:
+            policy_data.sel(**state_vals).variable.data.put(
+                0, res.x[0]
+            )  # will only work for scalar actions
+            value_data.sel(**state_vals).variable.data.put(
+                0, srv_function(pre_states, dr_best)
+            )
+        else:
+            print(f"Optimization failure at {state_vals}.")
+            print(root_res)
+
+            dr_best = {c: get_action_rule(root_res[i]) for i, c in enumerate(controls)}
+
+            policy_data.sel(**state_vals).variable.data.put(0, res.root)  # ?
+            value_data.sel(**state_vals).variable.data.put(
+                0, srv_function(pre_states, dr_best)
+            )
+
+    # use the xarray interpolator to create a decision rule.
+    dr_from_data = {
+        c: ar_from_data(
+            policy_data
+        )  # maybe needs to be more sensitive to the information set
+        for i, c in enumerate(controls)
+    }
+
+    dec_vf = block.get_decision_value_function(dr_from_data, continuation)
+    arr_vf = block.get_arrival_value_function(disc_params, dr_from_data, continuation)
+
+    return dr_from_data, dec_vf, arr_vf

HARK/model.py

@@ -31,12 +31,13 @@ class Control:
 
     Parameters
     ----------
-    args : list of str
+    iset : list of str
         The labels of the variables that are in the information set of this control.
     """
 
-    def __init__(self, args):
-        pass
+    def __init__(self, iset, upper_bound=None):
+        self.iset = iset
+        self.upper_bound = upper_bound
 
 
 def discretized_shock_dstn(shocks, disc_params):
@@ -161,11 +162,30 @@ def get_dynamics(self):
     def get_vars(self):
         return list(self.shocks.keys()) + list(self.dynamics.keys())
 
+    def get_controls(self):
+        """
+        TODO: Repeated in RBlock. Move to higher order class.
+        """
+        dyn = self.get_dynamics()
+
+        return [varn for varn in dyn if isinstance(dyn[varn], Control)]
+
     def transition(self, pre, dr):
         """
         Returns variable values given previous values and decision rule for all controls.
         """
-        return simulate_dynamics(self.dynamics, pre, dr)
+        dyn = self.dynamics.copy()
+
+        # don't simulate values that have already been given.
+        # this will break if there's a directly recursive label,
+        # i.e. if dynamics at time t for variable 'a'
+        # depend on state of 'a' at time t-1
+        # This is a forbidden case in CDC's design.
+        for varn in pre:
+            if varn in dyn:
+                del dyn[varn]
+
+        return simulate_dynamics(dyn, pre, dr)
 
     def calc_reward(self, vals):
         """

HARK/models/consumer.py

@@ -52,7 +52,7 @@
         "dynamics": {
             "b": lambda k, R, PermGroFac: k * R / PermGroFac,
             "m": lambda b, theta: b + theta,
-            "c": Control(["m"]),
+            "c": Control(["m"], upper_bound=lambda m: m),
             "a": lambda m, c: m - c,
         },
         "reward": {"u": lambda c, CRRA: c ** (1 - CRRA) / (1 - CRRA)},