getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting started with the Boltzmann Policy Distribution\n",
    "\n",
    "This notebook walks through how to calculate the Boltzmann policy distribution (BPD) for a simple environment and then use it to predict human behavior. You can use this as a tutorial to get started with using the Boltzmann policy distribution for your own environments!\n",
    "\n",
    "If you are running this notebook in Google Colab, it is recommended to use a GPU. You can enable GPU acceleration by going to **Runtime > Change runtime type** and selecting **GPU** from the dropdown."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, make sure you have installed the `boltzmann-policy-distribution` package, either from GitHub or PyPI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import bpd\n",
    "except ImportError:\n",
    "    %pip install boltzmann-policy-distribution ray[rllib]==1.11.0\n",
    "\n",
    "# Avoid warnings from RLlib:\n",
    "%pip install moviepy>=1.0.0 gputil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building an environment\n",
    "\n",
    "We'll build an environment similar to the one found in Figures 3 and 6 of [the paper](https://openreview.net/pdf?id=_l_QjPGN5ye) to demonstrate how to use the BPD. In this gridworld, the agent must repetitively travel to the an apple tree, pick an apple, and drop it in a basket. The world is in the shape of a ring with the apple tree and basket opposite one another, so the agent can travel either way around to pick apples. The agent receives 1 reward for dropping off an apple in the basket.\n",
    "\n",
    "Our implementation of the BPD is based on [RLlib](https://docs.ray.io/en/latest/rllib/index.html), which supports environments defined with the [OpenAI Gym](https://gym.openai.com/) interface, so our environment will use Gym."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Dict, Tuple\n",
    "from typing_extensions import TypedDict\n",
    "from ray.tune.registry import register_env\n",
    "\n",
    "import gym\n",
    "from gym import spaces\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class ApplePickingConfigDict(TypedDict):\n",
    "    ring_size: int\n",
    "    horizon: int\n",
    "\n",
    "\n",
    "DEFAULT_CONFIG: ApplePickingConfigDict = {\"ring_size\": 8, \"horizon\": 100}\n",
    "\n",
    "\n",
    "class ApplePickingEnv(gym.Env):\n",
    "    current_position: int\n",
    "    holding_apple: bool\n",
    "    timestep: int\n",
    "\n",
    "    def __init__(self, config: ApplePickingConfigDict = DEFAULT_CONFIG):\n",
    "        super().__init__()\n",
    "        self.config = {**DEFAULT_CONFIG, **config}\n",
    "\n",
    "        self.obs_space_size = self.config[\"ring_size\"] + 1\n",
    "        self.observation_space = spaces.Box(\n",
    "            low=np.zeros(self.obs_space_size, dtype=np.float32),\n",
    "            high=np.ones(self.obs_space_size, dtype=np.float32),\n",
    "        )\n",
    "        self.action_space = spaces.Discrete(2)\n",
    "\n",
    "    def _get_obs(self) -> np.ndarray:\n",
    "        obs = np.zeros(self.obs_space_size)\n",
    "        obs[self.current_position] = 1\n",
    "        if self.holding_apple:\n",
    "            obs[-1] = 1\n",
    "        return obs\n",
    "\n",
    "    def reset(self) -> np.ndarray:\n",
    "        self.current_position = 0\n",
    "        self.holding_apple = False\n",
    "        self.timestep = 0\n",
    "        return self._get_obs()\n",
    "\n",
    "    def step(self, action: int) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]:\n",
    "        if action == 0:\n",
    "            self.current_position += 1\n",
    "            # Wrap around the ring if necessary.\n",
    "            if self.current_position >= self.config[\"ring_size\"]:\n",
    "                self.current_position -= self.config[\"ring_size\"]\n",
    "        elif action == 1:\n",
    "            self.current_position -= 1\n",
    "            # Wrap around the ring if necessary.\n",
    "            if self.current_position < 0:\n",
    "                self.current_position += self.config[\"ring_size\"]\n",
    "        else:\n",
    "            assert False, \"Invalid action\"\n",
    "\n",
    "        reward = 0\n",
    "        # The basket is at position 0 in the ring.\n",
    "        if self.current_position == 0 and self.holding_apple:\n",
    "            # Drop off apple.\n",
    "            reward = 1\n",
    "            self.holding_apple = False\n",
    "        # The tree is halfway around the ring.\n",
    "        elif (\n",
    "            self.current_position == self.config[\"ring_size\"] // 2\n",
    "            and not self.holding_apple\n",
    "        ):\n",
    "            # Pick apple.\n",
    "            self.holding_apple = True\n",
    "\n",
    "        self.timestep += 1\n",
    "        done = self.timestep >= self.config[\"horizon\"]\n",
    "\n",
    "        return self._get_obs(), float(reward), done, {}\n",
    "\n",
    "\n",
    "register_env(\"apple_picking\", lambda env_config: ApplePickingEnv(env_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Human data\n",
    "\n",
    "The purpose of the BPD is to better predict human behavior, so let's generate some simulated human data to evaluate it with. We'll generate data for two consistent humans, one of whom always moves around the ring to the right and one of which always moves left. We can also generate data for a human which alternates going left and right (i.e., they are not as consistent)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from ray.rllib.evaluation import SampleBatch\n",
    "\n",
    "\n",
    "def generate_human_trajectory(action_sequence: List[int]) -> SampleBatch:\n",
    "    env = ApplePickingEnv()\n",
    "    done = False\n",
    "    obs_list = []\n",
    "    action_list = []\n",
    "    reward_list = []\n",
    "    obs = env.reset()\n",
    "    timestep = 0\n",
    "    while not done:\n",
    "        obs_list.append(obs)\n",
    "        action = action_sequence[timestep]\n",
    "        action_list.append(action)\n",
    "        obs, reward, done, info = env.step(action)\n",
    "        reward_list.append(reward)\n",
    "        timestep += 1\n",
    "    return SampleBatch(\n",
    "        {\n",
    "            SampleBatch.OBS: np.array(obs_list),\n",
    "            SampleBatch.ACTIONS: np.array(action_list),\n",
    "            SampleBatch.PREV_ACTIONS: np.array([0] + action_list[:-1]),\n",
    "            SampleBatch.REWARDS: np.array(reward_list),\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "human_trajectories = [\n",
    "    generate_human_trajectory([0] * 100),\n",
    "    generate_human_trajectory([1] * 100),\n",
    "    generate_human_trajectory(([0] * 8 + [1] * 8) * 7),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating a Boltzmann rational policy\n",
    "\n",
    "Before calculating the Boltzmann policy distribution for our environment, let's first calculate a Boltzmann rational policy to use as a baseline. We can do this with entropy-regularized policy gradient, where we add an entropy bonus to the policy gradient objective. We'll use RLlib's build-in [PPO](https://arxiv.org/abs/1707.06347) implementation to train the policy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "from ray.rllib.agents.ppo import PPOTrainer\n",
    "\n",
    "# Configure the algorithm.\n",
    "ppo_trainer = PPOTrainer(\n",
    "    config={\n",
    "        \"env\": \"apple_picking\",\n",
    "        \"framework\": \"torch\",\n",
    "        \"lr\": 1e-3,\n",
    "        \"gamma\": 0.9,\n",
    "        \"train_batch_size\": 2000,\n",
    "        \"sgd_minibatch_size\": 2000,\n",
    "        \"create_env_on_driver\": True,\n",
    "        # This adds the entropy regularizer to calculate the Boltzmann rational policy.\n",
    "        \"entropy_coeff\": 1,\n",
    "    }\n",
    ")\n",
    "\n",
    "for _ in tqdm.trange(10):\n",
    "    training_result = ppo_trainer.train()\n",
    "\n",
    "eval_result = ppo_trainer.evaluate()\n",
    "print(\"Mean reward = \", eval_result[\"evaluation\"][\"episode_reward_mean\"])\n",
    "ppo_trainer.cleanup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've trained a Boltzmann rational policy, let's see how well it predicts the simulated human actions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from ray.rllib.policy.torch_policy import TorchPolicy\n",
    "\n",
    "\n",
    "def evaluate_cross_entropy(policy: TorchPolicy, trajectory):\n",
    "    trajectory = trajectory.copy()\n",
    "    action_dist_inputs, _ = policy.model(\n",
    "        policy._lazy_tensor_dict(trajectory),\n",
    "        [\n",
    "            torch.tensor(state)[None].to(policy.device)\n",
    "            for state in policy.get_initial_state()\n",
    "        ],\n",
    "        np.array([len(trajectory)]),\n",
    "    )\n",
    "    action_distribution = policy.dist_class(action_dist_inputs)\n",
    "    return -action_distribution.logp(trajectory[SampleBatch.ACTIONS]).mean()\n",
    "\n",
    "\n",
    "boltzmann_rational_policy = ppo_trainer.get_policy()\n",
    "for traj_index, human_trajectory in enumerate(human_trajectories):\n",
    "    cross_entropy = evaluate_cross_entropy(boltzmann_rational_policy, human_trajectory)\n",
    "    print(\n",
    "        f\"Trajectory {traj_index}: Boltzmann rationality cross-entropy = {cross_entropy:.2f}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the Boltzmann rational policy does an okay job of predicting human behavior. The cross-entropy is lower than if it was predicting totally random actions (then it would be $\\ln(2) \\approx 0.7$). However, it can't take advantage of the fact that the first two human trajectories are consistent—that the human in those trajectories is always going around the ring the same way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the Boltzmann policy distribution (BPD)\n",
    "\n",
    "Now, let's calculate the Boltzmann policy distribution for this apple-picking environment. Unlike Boltzmann rationality, the BPD can take advantage of consistent human behavior to give better predictions of human actions.\n",
    "\n",
    "Unlike for Boltzmann rationality, we can't use built-in RLlib algorithms to calculate the BPD. Fortunately, we provide an RLlib-compatible BPD calculation algorithm! Using it requires just a bit more work than what we've already done.\n",
    "\n",
    "First, we have to implement a policy network that can produce a *distribution* over policies, rather than a single policy. We do this by conditioning on a latent vector which is passed in as part of the observation. We found that the best policy network architectures for the BPD use an attention mechanism to decide which parts of this latent vector to use for calculating action probabilities.\n",
    "\n",
    "The policy network also has to provide a *discriminator*, which is used in the implicit variational inference algorithm that calculates the BPD.\n",
    "\n",
    "The policy network implemented below looks complicated, but it's easy to modify for other environments. In fact, it should be possible to use it directly with other environments that only require fully-connected layers in their policy networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, List, Optional, Tuple, Union\n",
    "import gym\n",
    "from ray.rllib.models.catalog import ModelCatalog\n",
    "from ray.rllib.models.torch.torch_modelv2 import TorchModelV2\n",
    "from ray.rllib.policy.rnn_sequencing import add_time_dimension\n",
    "from ray.rllib.policy.sample_batch import SampleBatch\n",
    "import torch\n",
    "from torch import nn\n",
    "import numpy as np\n",
    "from ray.rllib.utils.typing import ModelConfigDict, TensorType\n",
    "\n",
    "\n",
    "class TransformerDiscriminator(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_dim: int,\n",
    "        out_dim: int,\n",
    "        size_hidden: int,\n",
    "        num_layers: int,\n",
    "        sum_over_seq: bool = True,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.sum_over_seq = sum_over_seq\n",
    "\n",
    "        self.encoder = nn.Linear(in_dim, size_hidden)\n",
    "        self.transformer = nn.TransformerEncoder(\n",
    "            nn.TransformerEncoderLayer(\n",
    "                d_model=size_hidden,\n",
    "                nhead=1,\n",
    "                dim_feedforward=size_hidden,\n",
    "                batch_first=True,\n",
    "            ),\n",
    "            num_layers=num_layers,\n",
    "        )\n",
    "        self.head = nn.Linear(size_hidden, out_dim)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        obs: torch.Tensor,\n",
    "        seq_lens: Union[torch.Tensor, np.ndarray],\n",
    "    ) -> torch.Tensor:\n",
    "        encoded_obs = self.encoder(obs)\n",
    "\n",
    "        if isinstance(seq_lens, np.ndarray):\n",
    "            seq_lens = torch.Tensor(seq_lens).int()\n",
    "        max_seq_len = encoded_obs.shape[0] // seq_lens.shape[0]\n",
    "        transformer_inputs = add_time_dimension(\n",
    "            encoded_obs,\n",
    "            max_seq_len=max_seq_len,\n",
    "            framework=\"torch\",\n",
    "            time_major=False,\n",
    "        )\n",
    "\n",
    "        transformer_outputs = self.transformer(transformer_inputs)\n",
    "        outputs: torch.Tensor = self.head(transformer_outputs)\n",
    "        if self.sum_over_seq:\n",
    "            outputs = outputs.sum(dim=1)\n",
    "        else:\n",
    "            outputs = outputs.reshape(-1, outputs.size()[-1])\n",
    "        return outputs\n",
    "\n",
    "\n",
    "class ApplePickingDistributionModel(TorchModelV2, nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        obs_space: gym.spaces.Space,\n",
    "        action_space: gym.spaces.Space,\n",
    "        num_outputs: int,\n",
    "        model_config: ModelConfigDict,\n",
    "        name: str,\n",
    "        latent_size: int,\n",
    "        num_heads: int = 4,\n",
    "        discriminate_sequences: bool = False,\n",
    "    ):\n",
    "        TorchModelV2.__init__(\n",
    "            self, obs_space, action_space, num_outputs, model_config, name\n",
    "        )\n",
    "        nn.Module.__init__(self)\n",
    "\n",
    "        assert self.model_config[\"vf_share_layers\"] is True\n",
    "        assert len(self.model_config[\"fcnet_hiddens\"]) == 3\n",
    "\n",
    "        self.num_heads = num_heads\n",
    "        self.latent_size = latent_size\n",
    "        self.discriminate_sequences = discriminate_sequences\n",
    "\n",
    "        in_dim = self.obs_space.shape[-1] + self.num_outputs\n",
    "        self.discriminator_net = self._build_discriminator(in_dim)\n",
    "        self.detached_discriminator_net = self._build_discriminator(in_dim)\n",
    "\n",
    "        self.backbone = nn.Sequential(\n",
    "            nn.Linear(\n",
    "                obs_space.shape[0] - latent_size, self.model_config[\"fcnet_hiddens\"][0]\n",
    "            ),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(\n",
    "                self.model_config[\"fcnet_hiddens\"][0],\n",
    "                self.model_config[\"fcnet_hiddens\"][1],\n",
    "            ),\n",
    "            nn.LeakyReLU(),\n",
    "        )\n",
    "        self.attention = nn.Linear(\n",
    "            self.model_config[\"fcnet_hiddens\"][1], self.num_heads * latent_size\n",
    "        )\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Linear(\n",
    "                self.model_config[\"fcnet_hiddens\"][1] + self.num_heads,\n",
    "                self.model_config[\"fcnet_hiddens\"][2],\n",
    "            ),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(self.model_config[\"fcnet_hiddens\"][2], num_outputs),\n",
    "        )\n",
    "\n",
    "        self.value_head = nn.Sequential(\n",
    "            nn.Linear(\n",
    "                self.model_config[\"fcnet_hiddens\"][1] + self.num_heads,\n",
    "                self.model_config[\"fcnet_hiddens\"][2],\n",
    "            ),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(self.model_config[\"fcnet_hiddens\"][2], 1),\n",
    "        )\n",
    "\n",
    "    def get_initial_state(self) -> List[np.ndarray]:\n",
    "        if self.discriminate_sequences:\n",
    "            return [np.zeros(1)]\n",
    "        else:\n",
    "            return super().get_initial_state()\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        input_dict: Dict[str, TensorType],\n",
    "        state: List[TensorType],\n",
    "        seq_lens: TensorType,\n",
    "    ) -> Tuple[TensorType, List[TensorType]]:\n",
    "        obs = input_dict[\"obs_flat\"].float()\n",
    "        self._last_flat_in = obs.reshape(obs.shape[0], -1)\n",
    "\n",
    "        self._features = self.backbone(self._last_flat_in[:, : -self.latent_size])\n",
    "        attention_weights = self.attention(self._features)\n",
    "        attention_weights = attention_weights.reshape(-1, 4, self.latent_size)\n",
    "        attention_weights = attention_weights.softmax(-1)\n",
    "\n",
    "        latent_attention_output = (\n",
    "            attention_weights * self._last_flat_in[:, None, -self.latent_size :]\n",
    "        ).sum(2)\n",
    "        head_input = torch.cat(\n",
    "            [self._features, latent_attention_output],\n",
    "            dim=1,\n",
    "        )\n",
    "\n",
    "        logits = self.head(head_input)\n",
    "        self._vf = self.value_head(head_input)[:, 0]\n",
    "\n",
    "        return logits, [s + 1 for s in state]\n",
    "\n",
    "    def value_function(self) -> TensorType:\n",
    "        return self._vf\n",
    "\n",
    "    def _build_discriminator(\n",
    "        self,\n",
    "        in_dim: int,\n",
    "        out_dim: int = 1,\n",
    "    ) -> nn.Module:\n",
    "        if self.discriminate_sequences:\n",
    "            return TransformerDiscriminator(\n",
    "                in_dim,\n",
    "                1,\n",
    "                self.model_config[\"fcnet_hiddens\"][0],\n",
    "                len(self.model_config[\"fcnet_hiddens\"]),\n",
    "            )\n",
    "        else:\n",
    "            dims = [in_dim] + self.model_config[\"fcnet_hiddens\"] + [out_dim]\n",
    "            layers: List[nn.Module] = []\n",
    "            for dim1, dim2 in zip(dims[:-1], dims[1:]):\n",
    "                layers.append(nn.Linear(dim1, dim2))\n",
    "                layers.append(nn.LeakyReLU())\n",
    "            layers = layers[:-1]\n",
    "            return nn.Sequential(*layers)\n",
    "\n",
    "    def discriminator(\n",
    "        self, input_dict, seq_lens: Optional[torch.Tensor] = None, detached=False\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Takes in a dictionary with observations and action probabilities and\n",
    "        outputs whether it thinks they came from this policy distribution or\n",
    "        the prior.\n",
    "\n",
    "        If detached is True, then this will run through a separate, \"detached\" copy of\n",
    "        the discriminator which will not propagate gradients to the main network.\n",
    "        \"\"\"\n",
    "\n",
    "        if detached:\n",
    "            self.detached_discriminator_net.load_state_dict(\n",
    "                self.discriminator_net.state_dict(keep_vars=False),\n",
    "            )\n",
    "            self.detached_discriminator_net.eval()\n",
    "            discriminator_net = self.detached_discriminator_net\n",
    "        else:\n",
    "            discriminator_net = self.discriminator_net\n",
    "\n",
    "        obs = input_dict[SampleBatch.OBS]\n",
    "        ac_probs = input_dict[SampleBatch.ACTION_PROB]\n",
    "        if not detached:\n",
    "            ac_probs = ac_probs + torch.normal(torch.zeros_like(ac_probs), 0.1)\n",
    "        net_input = torch.cat([obs, ac_probs], dim=1)\n",
    "        net_input[\n",
    "            :, self.obs_space.shape[-1] - self.latent_size : self.obs_space.shape[-1]\n",
    "        ] = 0\n",
    "\n",
    "        if self.discriminate_sequences:\n",
    "            return discriminator_net(net_input, seq_lens)\n",
    "        else:\n",
    "            return discriminator_net(net_input)\n",
    "\n",
    "\n",
    "ModelCatalog.register_custom_model(\n",
    "    \"apple_picking_distribution_model\",\n",
    "    ApplePickingDistributionModel,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've implemented the policy network, we're ready to calculate the BPD for this environment with implicit variational inference! We'll use similar code to when we trained the Boltzmann rational policy, but now we replace the PPO trainer with a BPD trainer. There are a few other important differences:\n",
    " * We need to wrap our environment in a `LatentEnvWrapper` from `bpd.envs.latent_wrapper`. This will pass in the latent vector which the BPD is conditioned to the policy network on as part of the observation.\n",
    " * We're using our custom model for the policy and value network, `ApplePickingDistributionModel`.\n",
    " * We pass in a temperature (the $1 / \\beta$ parameter from the paper), which controls how suboptimal the human is, and a prior concentration (the $\\alpha$ parameter from the paper), which controls how inconsistent the human is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "import torch\n",
    "from bpd.agents.bpd_trainer import BPDTrainer\n",
    "\n",
    "latent_size = 2\n",
    "\n",
    "# Configure the algorithm.\n",
    "bpd_trainer = BPDTrainer(\n",
    "    config={\n",
    "        \"env\": \"latent_wrapper\",\n",
    "        \"env_config\": {\n",
    "            \"env\": \"apple_picking\",\n",
    "            \"env_config\": {},\n",
    "            \"latent_dist\": lambda: np.random.normal(0, 1, 2),\n",
    "            \"episodes_per_latent\": 1,\n",
    "            \"agents_with_latent\": {0},\n",
    "        },\n",
    "        \"model\": {\n",
    "            \"custom_model\": \"apple_picking_distribution_model\",\n",
    "            \"max_seq_len\": 10,\n",
    "            \"vf_share_layers\": True,\n",
    "            \"custom_model_config\": {\n",
    "                \"latent_size\": latent_size,\n",
    "                \"discriminate_sequences\": True,\n",
    "            },\n",
    "            \"fcnet_hiddens\": [64, 64, 64],\n",
    "        },\n",
    "        \"temperature\": 0.1,\n",
    "        \"prior_concentration\": 0.2,\n",
    "        \"latent_size\": latent_size,\n",
    "        \"framework\": \"torch\",\n",
    "        \"lr\": 1e-3,\n",
    "        \"gamma\": 0.9,\n",
    "        \"num_sgd_iter\": 30,\n",
    "        \"train_batch_size\": 2000,\n",
    "        \"sgd_minibatch_size\": 2000,\n",
    "        \"num_gpus\": 1 if torch.cuda.is_available() else 0,\n",
    "        \"create_env_on_driver\": True,\n",
    "        \"vf_loss_coeff\": 0.01,\n",
    "    }\n",
    ")\n",
    "\n",
    "for _ in tqdm.trange(50):\n",
    "    training_result = bpd_trainer.train()\n",
    "\n",
    "eval_result = bpd_trainer.evaluate()\n",
    "print(\"Mean reward = \", eval_result[\"evaluation\"][\"episode_reward_mean\"])\n",
    "bpd_trainer.cleanup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Online prediction with the BPD\n",
    "\n",
    "Now that we've calculated the BPD for our environment, we can use it for online prediction of human behavior. Online prediction with the BPD requires solving an inference problem: we use the BPD as a prior over human policies and update it to a posterior after observing actions that the human takes. In the paper, we explored two ways of approximating this inference. The first explicitly models a posterior distribution over the latent vector with mean-field variational inference (MFVI). We have implemented this algorithm in the `VIActionPredictor` class. Try running the code below to see how the predictions made by the BPD compare to those made by Boltzmann rationality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bpd.latent_prediction import VIActionPredictor\n",
    "\n",
    "bpd_policy = bpd_trainer.get_policy()\n",
    "\n",
    "for traj_index, human_trajectory in enumerate(human_trajectories):\n",
    "    action_logits = VIActionPredictor(bpd_policy).predict_actions(human_trajectory)\n",
    "    action_distribution = bpd_policy.dist_class(action_logits)\n",
    "    cross_entropy = -action_distribution.logp(\n",
    "        human_trajectory[SampleBatch.ACTIONS]\n",
    "    ).mean()\n",
    "    print(\n",
    "        f\"Trajectory {traj_index}: Boltzmann policy distribution (using MFVI) cross-entropy = {cross_entropy:.2f}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how using the BPD results in much lower cross-entropy for predictions on the first two human trajectories, both of which are *consistent*. On the *inconsistent* third human trajectory, the BPD performs about the same as Boltzmann rationality. Fortunately, we find that real humans are usually consistent, so the BPD greatly outperforms Boltzmann rationality at predicting human behavior."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The other method we used to approximate the online prediction inference doesn't use an explicit representation of the posterior over the latent vector at all. Instead, it directly approximates the online prediction problem $p(a_t \\mid s_1, a_1, \\dots, s_t)$. To do this, we generate many trajectories from our calculated BPD policies and then train an autoregressive sequence model to predict the next action conditioned on previous states and actions. With enough trajectories, this approach should precisely approximate the exact online prediction inference.\n",
    "\n",
    "We have implemented an RLlib trainer called `DistillationPredictionTrainer` which can trains one policy to imitate another policy. In this case, we'll imitate our latent-conditioned BPD policy with an LSTM policy that *can't see the latent vector*. This will force the LSTM policy to learn to perform exactly the inference described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID\n",
    "from ray.rllib.policy.policy import PolicySpec\n",
    "from bpd.agents.distillation_prediction import DistillationPredictionTrainer\n",
    "\n",
    "PREDICTION_POLICY_ID = \"prediction\"\n",
    "env = ApplePickingEnv()\n",
    "prediction_trainer = DistillationPredictionTrainer(\n",
    "    config={\n",
    "        \"env\": \"latent_wrapper\",\n",
    "        \"env_config\": {\n",
    "            **bpd_trainer.config[\"env_config\"],\n",
    "        },\n",
    "        \"distillation_mapping_fn\": lambda policy_id: PREDICTION_POLICY_ID,\n",
    "        \"multiagent\": {\n",
    "            \"policies\": {\n",
    "                DEFAULT_POLICY_ID: PolicySpec(\n",
    "                    None,\n",
    "                    bpd_policy.observation_space,\n",
    "                    bpd_policy.action_space,\n",
    "                    {\"model\": bpd_trainer.config[\"model\"]},\n",
    "                ),\n",
    "                PREDICTION_POLICY_ID: PolicySpec(\n",
    "                    None,\n",
    "                    env.observation_space,\n",
    "                    env.action_space,\n",
    "                    {\n",
    "                        \"model\": {\n",
    "                            \"use_lstm\": True,\n",
    "                            \"lstm_cell_size\": 32,\n",
    "                            \"lstm_use_prev_action\": True,\n",
    "                            \"max_seq_len\": 100,\n",
    "                        },\n",
    "                    },\n",
    "                ),\n",
    "            },\n",
    "            \"policies_to_train\": [PREDICTION_POLICY_ID],\n",
    "            \"policy_mapping_fn\": lambda agent_id, episodes, worker, **kwargs: DEFAULT_POLICY_ID,\n",
    "        },\n",
    "        \"framework\": \"torch\",\n",
    "        \"lr\": 1e-3,\n",
    "        \"num_sgd_iter\": 3,\n",
    "        \"train_batch_size\": 500,\n",
    "        \"sgd_minibatch_size\": 500,\n",
    "        \"num_gpus\": 1 if torch.cuda.is_available() else 0,\n",
    "        \"create_env_on_driver\": True,\n",
    "        \"num_workers\": 0,\n",
    "    }\n",
    ")\n",
    "prediction_trainer.set_weights(bpd_trainer.get_weights())\n",
    "\n",
    "for _ in tqdm.trange(40):\n",
    "    training_result = prediction_trainer.train()\n",
    "\n",
    "print(\n",
    "    \"Train cross entropy = \",\n",
    "    training_result[\"info\"][\"learner\"][PREDICTION_POLICY_ID][\"learner_stats\"][\n",
    "        \"cross_entropy\"\n",
    "    ],\n",
    ")\n",
    "prediction_trainer.cleanup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how our LSTM prediction policy does on the simulated human data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bpd_prediction_policy = prediction_trainer.get_policy(PREDICTION_POLICY_ID)\n",
    "for traj_index, human_trajectory in enumerate(human_trajectories):\n",
    "    cross_entropy = evaluate_cross_entropy(bpd_prediction_policy, human_trajectory)\n",
    "    print(f\"Trajectory {traj_index}: BPD cross-entropy = {cross_entropy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like with the MFVI inference method, the BPD outperforms Boltzmann rationality for the first two trajectories where the simulated human is consistent. Notice how fast the LSTM prediction is—much faster than MFVI. This makes it useful for applications where we need to predict human behavior in real time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you should be ready to use the Boltzmann policy distribution for modeling humans in new environments. If you want to see how we used the BPD for a more complex environment, Overcooked, see the [README](https://github.com/cassidylaidlaw/boltzmann-policy-distribution/blob/main/README.md). Feel free to reach out with any questions to [cassidy_laidlaw@berkeley.edu](mailto:cassidy_laidlaw@berkeley.edu).\n",
    "\n",
    "Happy human modeling!"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "9d92d321245a4a8fcb6f4caced66fff48020fd1dc76d949a2f6e29bd752cc591"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('rllib')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}