lending.py

# coding=utf-8
# Copyright 2022 The ML Fairness Gym Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python2, python3
"""ML Fairness gym loan environment.

This environment is meant to be a hello-world example to the gym as well as
serve as a template for writing future environments.

In each step step, the agent decides whether to accept or reject an application.

Applicant features are generated by a mixture model which also determines the
likelihood of defaulting.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import enum
from typing import List, Optional

from absl import logging
import attr
import core
from environments import lending_params
from spaces import multinomial
from gym import spaces
import matplotlib
import matplotlib.pyplot as plt
import numpy as np


# Used for rending applicant features.
_MARKERS = matplotlib.markers.MarkerStyle.filled_markers


class LoanDecision(enum.IntEnum):
  """Enum representing possible loan decisions."""
  REJECT = 0
  ACCEPT = 1


class _CashUpdater(core.StateUpdater):
  """Changes bank_cash as a result of an action."""

  def update(self, state, action):
    params = state.params
    if action == LoanDecision.REJECT:
      return
    if state.will_default:
      state.bank_cash -= params.loan_amount
    else:
      state.bank_cash += params.loan_amount * params.interest_rate


class _ApplicantSampler(core.StateUpdater):
  """Samples a new applicant."""

  def update(self, state, action):
    del action  # Unused.
    params = state.params
    new_applicant = params.applicant_distribution.sample(state.rng)
    state.applicant_features = np.clip(new_applicant.features,
                                       params.min_observation,
                                       params.max_observation)
    state.group = new_applicant.group
    state.group_id = np.argmax(new_applicant.group)
    state.will_default = new_applicant.will_default


@attr.s(cmp=False)  # Use core.State's equality methods.
class State(core.State):
  """State object for lending environments."""

  # Random number generator for the simulation.
  rng = attr.ib()  # type: np.random.RandomState

  # State parameters that can evolve over time.
  params = attr.ib()  # type: lending_params.Params

  # Number of loans available for the bank.
  bank_cash = attr.ib()  # type: float

  # Applicant-related attributes are Optional with defaults of None so that the
  # object can be initialized in two steps, first with applicant attributes as
  # None, then a StateUpdater is used to fill in the applicant features.
  applicant_features = attr.ib(default=None)  # type: Optional[np.ndarray]
  group = attr.ib(default=None)  # type: Optional[List[int]]
  group_id = attr.ib(default=None)  # type: Optional[int]
  will_default = attr.ib(default=None)  # type: Optional[bool]


class BaseLendingEnv(core.FairnessEnv):
  """Base loan decision environment.

  In each step, the agent decides whether to accept or reject an
  application.

  The base class is abstract.
  """

  metadata = {'render.modes': ['human']}
  default_param_builder = lending_params.Params
  group_membership_var = 'group'
  _cash_updater = _CashUpdater()
  _parameter_updater = core.NoUpdate()
  _applicant_updater = _ApplicantSampler()

  def __init__(self, params = None):
    params = (
        self.default_param_builder() if params is None else params
        )  # type: lending_params.Params

    # The action space of the agent is Accept/Reject.
    self.action_space = spaces.Discrete(2)

    # Bank's cash is a scalar and cannot be negative.
    bank_cash_space = spaces.Box(
        low=0, high=params.max_cash, shape=(), dtype=np.float32)

    # Two-dimensional observation space describes each loan applicant.
    loan_applicant_space = spaces.Box(
        params.min_observation,
        params.max_observation,
        dtype=np.float32,
        shape=(params.applicant_distribution.dim,))

    group_space = spaces.MultiBinary(params.num_groups)

    self.observable_state_vars = {
        'bank_cash': bank_cash_space,
        'applicant_features': loan_applicant_space,
        'group': group_space
    }

    super(BaseLendingEnv, self).__init__(params)
    self._state_init()

  def _state_init(self, rng=None):
    self.state = State(
        # Copy in case state.params get mutated, initial_params stays pristine.
        params=copy.deepcopy(self.initial_params),
        rng=rng or np.random.RandomState(),
        bank_cash=self.initial_params.bank_starting_cash)
    self._applicant_updater.update(self.state, None)

  def reset(self):
    """Resets the environment."""
    self._state_init(self.state.rng)
    return super(BaseLendingEnv, self).reset()

  def _is_done(self):
    """Returns True if the bank cash is less than loan_amount."""
    return self.state.bank_cash < self.state.params.loan_amount

  def _step_impl(self, state, action):
    """Run one timestep of the environment's dynamics.

    In a single step, the agent decides whether to accept or reject an
    application.

    The potential payoffs of rejected application are always 0.
    If an application is accepted, the payoffs are:
      -loan_amount if the applicant defaults.
      +loan_amount*interest_rate if the applicant successfully pays back.

    Args:
      state: A `State` object containing the current state.
      action: An action in `action_space`.

    Returns:
      A `State` object containing the updated state.
    """

    self._cash_updater.update(self.state, action)
    self._parameter_updater.update(self.state, action)
    self._applicant_updater.update(self.state, action)
    return self.state

  def render(self, mode='human'):
    """Renders the history and current state using matplotlib.

    Args:
      mode: string indicating the rendering mode. The only supported mode is
        `human`.
    """

    if mode == 'human':
      if self.state.params.applicant_distribution.dim != 2:
        raise NotImplementedError(
            'Cannot render if applicant features are not exactly 2 dimensional. '
            'Got %d dimensional applicant features.' %
            self.state.params.applicant_distribution.dim)

      plt.figure(figsize=(12, 4))
      plt.subplot(1, 2, 1)
      plt.xlim(-2, 2)
      plt.ylim(-2, 2)
      plt.title('Applicant Features')
      plt.xticks([], [])
      plt.yticks([], [])
      for state, action in self.history:
        if action == 1:
          x, y = state.applicant_features
          color = 'r' if state.will_default else 'b'
          plt.plot([x], [y], _MARKERS[state.group_id] + color, markersize=12)
      plt.xlabel('Feature 1')
      plt.ylabel('Feature 2')

      x, y = self.state.applicant_features

      plt.plot([x], [y], _MARKERS[self.state.group_id] + 'k', markersize=15)

      plt.subplot(1, 2, 2)
      plt.title('Cash')
      plt.plot([state.bank_cash for state, _ in self.history] +
               [self.state.bank_cash])
      plt.ylabel('# loans available')
      plt.xlabel('Time')
      plt.tight_layout()
    else:
      super(BaseLendingEnv, self).render(mode)  # Raises NotImplementedError


class SimpleLoans(BaseLendingEnv):
  """Simple lending environment.

  Applicants have 2D features which can be used to determine whether they have
  high or low likelihood of success.
  """
  default_param_builder = lending_params.Params


class DifferentialExpressionEnv(BaseLendingEnv):
  """Lending environment with groups that present creditworthiness differently.

  Applicants have 2D features which can be used to determine whether they have
  high or low likelihood of success, but the mapping is different for the
  different groups.
  """
  default_param_builder = lending_params.DifferentialExpressionParams


class _CreditShift(core.StateUpdater):
  """Updates the cluster probabilities based on the repayment."""

  def update(self, state, action):
    """Updates the cluster probabilities based on the repayment.

    Successful repayment raises one's credit score and default lowers one's
    credit score. This is expressed by moving a small amount of probability mass
    (representing an individual) from one credit-score cluster to an adjacent
    one.

    This change in credit only happens if the applicant is accepted. Rejected
    applicants experience no change in their score.

    state.params is mutated in place; nothing is returned.

    Args:
      state: A core.State object.
      action: a `LoanDecision`.
    """

    if action == LoanDecision.REJECT:
      return

    params = state.params
    group_id = state.group_id

    # Group should always be a one-hot encoding of group_id. This assert tests
    # that these two values have not somehow gotten out of sync.
    assert state.group_id == np.argmax(
        state.group), 'Group id %s. group %s' % (state.group_id,
                                                 np.argmax(state.group))

    # Cast to list so we can mutate it.
    cluster_probs = list(
        params.applicant_distribution.components[group_id].weights)

    rng = np.random.RandomState()
    for _ in range(10):
      group = params.applicant_distribution.components[group_id].sample(
          rng).group
      assert np.array_equal(group, state.group), (
          'Sampling from the component that is indexed here does not give '
          'members of the group that is intended to be affected. Something is '
          'quite wrong. Check that your group ids are in order in the credit '
          'cluster spec. sampled group_id %s vs state.group %s. '
          'Component[%d]: %s' %
          (group, state.group, group_id,
           params.applicant_distribution.components[group_id]))

    # Assert argmax gives the right index.
    for idx, component in enumerate(
        params.applicant_distribution.components[group_id].components):
      credit_score = component.sample(rng).features
      assert np.argmax(credit_score) == idx, '%s vs %s' % (credit_score, idx)

    # This applicant has their credit score lowered or raised.
    cluster_id = np.argmax(state.applicant_features)
    new_cluster = (cluster_id - 1 if state.will_default else cluster_id + 1)

    # Prevents falling off the edges of the cluster array.
    new_cluster = min(new_cluster, len(cluster_probs) - 1)
    new_cluster = max(new_cluster, 0)

    # Prevents moving more probability mass than this bucket has.
    assert cluster_probs[cluster_id] > 0, (
        'This cluster was sampled but has no mass. %d. Full distribution %s' %
        (cluster_id, cluster_probs))

    mass_to_shift = min(params.cluster_shift_increment,
                        cluster_probs[cluster_id])

    # Mutates params.cluster_probs[group_id].
    cluster_probs[cluster_id] -= mass_to_shift
    cluster_probs[new_cluster] += mass_to_shift
    logging.debug('Group %d: Moving mass %f from %d to %d', group_id,
                  mass_to_shift, cluster_id, new_cluster)

    assert np.abs(np.sum(cluster_probs) -
                  1) < 1e-6, 'Cluster probs must sum to 1.'
    assert all([prob >= 0 for prob in cluster_probs
               ]), 'Cluster probs must be non-negative'

    state.params.applicant_distribution.components[
        group_id].weights = cluster_probs


class DelayedImpactEnv(BaseLendingEnv):
  """Lending environment in which outcomes affect future credit.

  Each applicant has a credit score which causally determines their likelihood
  of success. Applicants who default have their credit lowered while applicants
  who pay back have their credit raised.

  Based on the environment described in Liu et al's Delayed Impact of Machine
  Learning: https://arxiv.org/abs/1803.04383
  """
  default_param_builder = lending_params.DelayedImpactParams
  _parameter_updater = _CreditShift()

  def __init__(self, params=None):
    super(DelayedImpactEnv, self).__init__(params)
    self.observable_state_vars['applicant_features'] = multinomial.Multinomial(
        self.initial_params.applicant_distribution.dim, 1)
    self.observation_space = spaces.Dict(self.observable_state_vars)