examples/train_meta_standalone.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "421aea2b-85b8-4a3e-8673-150ee15eb9a3",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/corl-team/xland-minigrid/blob/main/examples/train_meta_standalone.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08db9383-e3e3-423f-b729-af3b8a310805",
   "metadata": {},
   "source": [
    "# Meta-task PPO "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcda9a2f-1685-4962-97a7-3d79754195d6",
   "metadata": {},
   "source": [
    "> ⚠️ Ensure you select a GPU from `Runtime > Change runtime type`. ⚠️\n",
    "\n",
    "> 🔥 Instances with multiple T4 gpus are available on Kaggle for free! Multi-gpu can speed up training with `pmap`. 🔥"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3446372-0dcb-45bc-9539-7da52b43b0bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# jax is already installed on the colab, uncomment only if needed\n",
    "# !pip install --upgrade \"jax[cuda11_pip]\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
    "\n",
    "# !pip install 'xminigrid[baselines]'\n",
    "!pip install \"xminigrid[baselines] @ git+https://github.com/corl-team/xland-minigrid.git\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2df3e354-ce7f-472c-bbb6-05f1b88e343f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import math\n",
    "from typing import TypedDict\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax.tree_util as jtu\n",
    "import flax\n",
    "import flax.linen as nn\n",
    "import distrax\n",
    "import optax\n",
    "import imageio\n",
    "import wandb\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from flax import struct\n",
    "from flax.linen.initializers import glorot_normal, orthogonal, zeros_init\n",
    "from flax.training.train_state import TrainState\n",
    "from flax.jax_utils import replicate, unreplicate\n",
    "from dataclasses import asdict, dataclass\n",
    "from functools import partial\n",
    "\n",
    "import xminigrid\n",
    "from xminigrid.environment import Environment, EnvParams\n",
    "from xminigrid.wrappers import GymAutoResetWrapper\n",
    "from xminigrid.benchmarks import Benchmark"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae8cbcf7-4008-440c-9ebf-0e7d26711253",
   "metadata": {},
   "source": [
    "## Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bcd8fd6-d359-4099-97e7-ad7dd850b2bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model adapted from minigrid baselines:\n",
    "# https://github.com/lcswillems/rl-starter-files/blob/master/model.py\n",
    "\n",
    "# custom RNN cell, which is more convenient that default in flax\n",
    "class GRU(nn.Module):\n",
    "    hidden_dim: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, xs, init_state):\n",
    "        seq_len, input_dim = xs.shape\n",
    "        # this init might not be optimal, for example bias for reset gate should be -1 (for now ok)\n",
    "        Wi = self.param(\"Wi\", glorot_normal(in_axis=1, out_axis=0), (self.hidden_dim * 3, input_dim))\n",
    "        Wh = self.param(\"Wh\", orthogonal(column_axis=0), (self.hidden_dim * 3, self.hidden_dim))\n",
    "        bi = self.param(\"bi\", zeros_init(), (self.hidden_dim * 3,))\n",
    "        bn = self.param(\"bn\", zeros_init(), (self.hidden_dim,))\n",
    "\n",
    "        def _step_fn(h, x):\n",
    "            igates = jnp.split(Wi @ x + bi, 3)\n",
    "            hgates = jnp.split(Wh @ h, 3)\n",
    "\n",
    "            reset = nn.sigmoid(igates[0] + hgates[0])\n",
    "            update = nn.sigmoid(igates[1] + hgates[1])\n",
    "            new = nn.tanh(igates[2] + reset * (hgates[2] + bn))\n",
    "            next_h = (1 - update) * new + update * h\n",
    "\n",
    "            return next_h, next_h\n",
    "\n",
    "        last_state, all_states = jax.lax.scan(_step_fn, init=init_state, xs=xs)\n",
    "        return all_states, last_state\n",
    "\n",
    "class RNNModel(nn.Module):\n",
    "    hidden_dim: int\n",
    "    num_layers: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, xs, init_state):\n",
    "        # xs: [seq_len, input_dim]\n",
    "        # init_state: [num_layers, hidden_dim]\n",
    "        outs, states = [], []\n",
    "        for layer in range(self.num_layers):\n",
    "            xs, state = GRU(hidden_dim=self.hidden_dim)(xs, init_state[layer])\n",
    "            outs.append(xs)\n",
    "            states.append(state)\n",
    "\n",
    "        # sum outputs from all layers, kinda like in ResNet\n",
    "        return jnp.array(outs).sum(0), jnp.array(states)\n",
    "\n",
    "BatchedRNNModel = flax.linen.vmap(\n",
    "    RNNModel, variable_axes={\"params\": None}, split_rngs={\"params\": False}, axis_name=\"batch\"\n",
    ")\n",
    "\n",
    "\n",
    "class ActorCriticInput(TypedDict):\n",
    "    observation: jax.Array\n",
    "    prev_action: jax.Array\n",
    "    prev_reward: jax.Array\n",
    "\n",
    "\n",
    "class ActorCriticRNN(nn.Module):\n",
    "    num_actions: int\n",
    "    action_emb_dim: int = 16\n",
    "    rnn_hidden_dim: int = 64\n",
    "    rnn_num_layers: int = 1\n",
    "    head_hidden_dim: int = 64\n",
    "    img_obs: bool = False\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, inputs: ActorCriticInput, hidden: jax.Array) -> tuple[distrax.Categorical, jax.Array, jax.Array]:\n",
    "        B, S = inputs[\"observation\"].shape[:2]\n",
    "        # encoder from https://github.com/lcswillems/rl-starter-files/blob/master/model.py\n",
    "        if self.img_obs:\n",
    "            img_encoder = nn.Sequential(\n",
    "                [\n",
    "                    nn.Conv(16, (3, 3), strides=2, padding=\"VALID\", kernel_init=orthogonal(math.sqrt(2))),\n",
    "                    nn.relu,\n",
    "                    nn.Conv(32, (3, 3), strides=2, padding=\"VALID\", kernel_init=orthogonal(math.sqrt(2))),\n",
    "                    nn.relu,\n",
    "                    nn.Conv(32, (3, 3), strides=2, padding=\"VALID\", kernel_init=orthogonal(math.sqrt(2))),\n",
    "                    nn.relu,\n",
    "                    nn.Conv(32, (3, 3), strides=2, padding=\"VALID\", kernel_init=orthogonal(math.sqrt(2))),\n",
    "                ]\n",
    "            )\n",
    "        else:\n",
    "            img_encoder = nn.Sequential(\n",
    "                [\n",
    "                    nn.Conv(16, (2, 2), padding=\"VALID\", kernel_init=orthogonal(math.sqrt(2))),\n",
    "                    nn.relu,\n",
    "                    nn.Conv(32, (2, 2), padding=\"VALID\", kernel_init=orthogonal(math.sqrt(2))),\n",
    "                    nn.relu,\n",
    "                    nn.Conv(64, (2, 2), padding=\"VALID\", kernel_init=orthogonal(math.sqrt(2))),\n",
    "                    nn.relu,\n",
    "                ]\n",
    "            )\n",
    "        action_encoder = nn.Embed(self.num_actions, self.action_emb_dim)\n",
    "\n",
    "        rnn_core = BatchedRNNModel(self.rnn_hidden_dim, self.rnn_num_layers)\n",
    "        actor = nn.Sequential(\n",
    "            [\n",
    "                nn.Dense(self.head_hidden_dim, kernel_init=orthogonal(2)),\n",
    "                nn.tanh,\n",
    "                nn.Dense(self.num_actions, kernel_init=orthogonal(0.01)),\n",
    "            ]\n",
    "        )\n",
    "        critic = nn.Sequential(\n",
    "            [\n",
    "                nn.Dense(self.head_hidden_dim, kernel_init=orthogonal(2)),\n",
    "                nn.tanh,\n",
    "                nn.Dense(1, kernel_init=orthogonal(1.0)),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        # [batch_size, seq_len, ...]\n",
    "        obs_emb = img_encoder(inputs[\"observation\"]).reshape(B, S, -1)\n",
    "        act_emb = action_encoder(inputs[\"prev_action\"])\n",
    "        # [batch_size, seq_len, hidden_dim + act_emb_dim + 1]\n",
    "        out = jnp.concatenate([obs_emb, act_emb, inputs[\"prev_reward\"][..., None]], axis=-1)\n",
    "        # core networks\n",
    "        out, new_hidden = rnn_core(out, hidden)\n",
    "        dist = distrax.Categorical(logits=actor(out))\n",
    "        values = critic(out)\n",
    "\n",
    "        return dist, jnp.squeeze(values, axis=-1), new_hidden\n",
    "\n",
    "    def initialize_carry(self, batch_size):\n",
    "        return jnp.zeros((batch_size, self.rnn_num_layers, self.rnn_hidden_dim))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa17fc29-d932-49df-bc3b-4c85bb4d4fce",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e326b09-6289-4744-84fa-75f134f7979a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training stuff\n",
    "class Transition(struct.PyTreeNode):\n",
    "    done: jax.Array\n",
    "    action: jax.Array\n",
    "    value: jax.Array\n",
    "    reward: jax.Array\n",
    "    log_prob: jax.Array\n",
    "    obs: jax.Array\n",
    "    # for rnn policy\n",
    "    prev_action: jax.Array\n",
    "    prev_reward: jax.Array\n",
    "\n",
    "\n",
    "def calculate_gae(\n",
    "    transitions: Transition,\n",
    "    last_val: jax.Array,\n",
    "    gamma: float,\n",
    "    gae_lambda: float,\n",
    ") -> tuple[jax.Array, jax.Array]:\n",
    "    # single iteration for the loop\n",
    "    def _get_advantages(gae_and_next_value, transition):\n",
    "        gae, next_value = gae_and_next_value\n",
    "        delta = transition.reward + gamma * next_value * (1 - transition.done) - transition.value\n",
    "        gae = delta + gamma * gae_lambda * (1 - transition.done) * gae\n",
    "        return (gae, transition.value), gae\n",
    "\n",
    "    _, advantages = jax.lax.scan(\n",
    "        _get_advantages,\n",
    "        (jnp.zeros_like(last_val), last_val),\n",
    "        transitions,\n",
    "        reverse=True,\n",
    "    )\n",
    "    # advantages and values (Q)\n",
    "    return advantages, advantages + transitions.value\n",
    "\n",
    "\n",
    "def ppo_update_networks(\n",
    "    train_state: TrainState,\n",
    "    transitions: Transition,\n",
    "    init_hstate: jax.Array,\n",
    "    advantages: jax.Array,\n",
    "    targets: jax.Array,\n",
    "    clip_eps: float,\n",
    "    vf_coef: float,\n",
    "    ent_coef: float,\n",
    "):\n",
    "    # NORMALIZE ADVANTAGES\n",
    "    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)\n",
    "\n",
    "    def _loss_fn(params):\n",
    "        # RERUN NETWORK\n",
    "        dist, value, _ = train_state.apply_fn(\n",
    "            params,\n",
    "            {\n",
    "                # [batch_size, seq_len, ...]\n",
    "                \"observation\": transitions.obs,\n",
    "                \"prev_action\": transitions.prev_action,\n",
    "                \"prev_reward\": transitions.prev_reward,\n",
    "            },\n",
    "            init_hstate,\n",
    "        )\n",
    "        log_prob = dist.log_prob(transitions.action)\n",
    "\n",
    "        # CALCULATE VALUE LOSS\n",
    "        value_pred_clipped = transitions.value + (value - transitions.value).clip(-clip_eps, clip_eps)\n",
    "        value_loss = jnp.square(value - targets)\n",
    "        value_loss_clipped = jnp.square(value_pred_clipped - targets)\n",
    "        value_loss = 0.5 * jnp.maximum(value_loss, value_loss_clipped).mean()\n",
    "\n",
    "        # CALCULATE ACTOR LOSS\n",
    "        ratio = jnp.exp(log_prob - transitions.log_prob)\n",
    "        actor_loss1 = advantages * ratio\n",
    "        actor_loss2 = advantages * jnp.clip(ratio, 1.0 - clip_eps, 1.0 + clip_eps)\n",
    "        actor_loss = -jnp.minimum(actor_loss1, actor_loss2).mean()\n",
    "        entropy = dist.entropy().mean()\n",
    "\n",
    "        total_loss = actor_loss + vf_coef * value_loss - ent_coef * entropy\n",
    "        return total_loss, (value_loss, actor_loss, entropy)\n",
    "\n",
    "    (loss, (vloss, aloss, entropy)), grads = jax.value_and_grad(_loss_fn, has_aux=True)(train_state.params)\n",
    "    (loss, vloss, aloss, entropy, grads) = jax.lax.pmean((loss, vloss, aloss, entropy, grads), axis_name=\"devices\")\n",
    "    train_state = train_state.apply_gradients(grads=grads)\n",
    "    update_info = {\n",
    "        \"total_loss\": loss,\n",
    "        \"value_loss\": vloss,\n",
    "        \"actor_loss\": aloss,\n",
    "        \"entropy\": entropy,\n",
    "    }\n",
    "    return train_state, update_info\n",
    "\n",
    "\n",
    "# for evaluation (evaluate for N consecutive episodes, sum rewards)\n",
    "# N=1 single task, N>1 for meta-RL\n",
    "class RolloutStats(struct.PyTreeNode):\n",
    "    reward: jax.Array = jnp.asarray(0.0)\n",
    "    length: jax.Array = jnp.asarray(0)\n",
    "    episodes: jax.Array = jnp.asarray(0)\n",
    "\n",
    "\n",
    "def rollout(\n",
    "    rng: jax.Array,\n",
    "    env: Environment,\n",
    "    env_params: EnvParams,\n",
    "    train_state: TrainState,\n",
    "    init_hstate: jax.Array,\n",
    "    num_consecutive_episodes: int = 1,\n",
    ") -> RolloutStats:\n",
    "    def _cond_fn(carry):\n",
    "        rng, stats, timestep, prev_action, prev_reward, hstate = carry\n",
    "        return jnp.less(stats.episodes, num_consecutive_episodes)\n",
    "\n",
    "    def _body_fn(carry):\n",
    "        rng, stats, timestep, prev_action, prev_reward, hstate = carry\n",
    "\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "        dist, _, hstate = train_state.apply_fn(\n",
    "            train_state.params,\n",
    "            {\n",
    "                \"observation\": timestep.observation[None, None, ...],\n",
    "                \"prev_action\": prev_action[None, None, ...],\n",
    "                \"prev_reward\": prev_reward[None, None, ...],\n",
    "            },\n",
    "            hstate,\n",
    "        )\n",
    "        action = dist.sample(seed=_rng).squeeze()\n",
    "        timestep = env.step(env_params, timestep, action)\n",
    "\n",
    "        stats = stats.replace(\n",
    "            reward=stats.reward + timestep.reward,\n",
    "            length=stats.length + 1,\n",
    "            episodes=stats.episodes + timestep.last(),\n",
    "        )\n",
    "        carry = (rng, stats, timestep, action, timestep.reward, hstate)\n",
    "        return carry\n",
    "\n",
    "    timestep = env.reset(env_params, rng)\n",
    "    prev_action = jnp.asarray(0)\n",
    "    prev_reward = jnp.asarray(0)\n",
    "    init_carry = (rng, RolloutStats(), timestep, prev_action, prev_reward, init_hstate)\n",
    "\n",
    "    final_carry = jax.lax.while_loop(_cond_fn, _body_fn, init_val=init_carry)\n",
    "    return final_carry[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aa618e1-03a3-4686-888e-ca771d3b67d1",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "407f5829-b77d-4fb5-b12e-43c152b18d06",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class TrainConfig:\n",
    "    env_id: str = \"XLand-MiniGrid-R1-8x8\"\n",
    "    benchmark_id: str = \"trivial-1m\"\n",
    "    img_obs: bool = False\n",
    "    # agent\n",
    "    action_emb_dim: int = 16\n",
    "    rnn_hidden_dim: int = 64\n",
    "    rnn_num_layers: int = 1\n",
    "    head_hidden_dim: int = 64\n",
    "    # training\n",
    "    num_envs: int = 1024\n",
    "    num_steps_per_env: int = 4096\n",
    "    num_steps_per_update: int = 32\n",
    "    update_epochs: int = 1\n",
    "    num_minibatches: int = 16\n",
    "    total_timesteps: int = 100_000_000\n",
    "    lr: float = 0.001\n",
    "    clip_eps: float = 0.2\n",
    "    gamma: float = 0.99\n",
    "    gae_lambda: float = 0.95\n",
    "    ent_coef: float = 0.01\n",
    "    vf_coef: float = 0.5\n",
    "    max_grad_norm: float = 0.5\n",
    "    eval_num_envs: int = 256\n",
    "    eval_num_episodes: int = 10\n",
    "    eval_seed: int = 42\n",
    "    train_seed: int = 42\n",
    "\n",
    "    def __post_init__(self):\n",
    "        num_devices = jax.local_device_count()\n",
    "        \n",
    "        # splitting computation across all available devices\n",
    "        self.num_envs_per_device = self.num_envs // num_devices\n",
    "        self.total_timesteps_per_device = self.total_timesteps // num_devices\n",
    "        self.eval_num_envs_per_device = self.eval_num_envs // num_devices\n",
    "        assert self.num_envs % num_devices == 0\n",
    "        \n",
    "        self.num_meta_updates = round(self.total_timesteps_per_device / (self.num_envs_per_device * self.num_steps_per_env))\n",
    "        self.num_inner_updates = self.num_steps_per_env // self.num_steps_per_update\n",
    "        assert self.num_steps_per_env % self.num_steps_per_update == 0\n",
    "        print(f\"Num devices: {num_devices}, Num meta updates: {self.num_meta_updates}\")\n",
    "\n",
    "\n",
    "def make_states(config: TrainConfig):\n",
    "    # for learning rage scheduling\n",
    "    def linear_schedule(count):\n",
    "        total_inner_updates = config.num_minibatches * config.update_epochs * config.num_inner_updates\n",
    "        frac = 1.0 - (count // total_inner_updates) / config.num_meta_updates\n",
    "        return config.lr * frac\n",
    "\n",
    "    # setup environment\n",
    "    if \"XLand\" not in config.env_id:\n",
    "        raise ValueError(\"Only meta-task environments are supported.\")\n",
    "\n",
    "    env, env_params = xminigrid.make(config.env_id)\n",
    "    env = GymAutoResetWrapper(env)\n",
    "\n",
    "    # enabling image observations if needed\n",
    "    if config.img_obs:\n",
    "        from xminigrid.experimental.img_obs import RGBImgObservationWrapper\n",
    "\n",
    "        env = RGBImgObservationWrapper(env)\n",
    "    \n",
    "    # loading benchmark\n",
    "    benchmark = xminigrid.load_benchmark(config.benchmark_id)\n",
    "\n",
    "    # set up training state\n",
    "    rng = jax.random.key(config.train_seed)\n",
    "    rng, _rng = jax.random.split(rng)\n",
    "\n",
    "    network = ActorCriticRNN(\n",
    "        num_actions=env.num_actions(env_params),\n",
    "        action_emb_dim=config.action_emb_dim,\n",
    "        rnn_hidden_dim=config.rnn_hidden_dim,\n",
    "        rnn_num_layers=config.rnn_num_layers,\n",
    "        head_hidden_dim=config.head_hidden_dim,\n",
    "        img_obs=config.img_obs,\n",
    "    )\n",
    "    # [batch_size, seq_len, ...]\n",
    "    init_obs = {\n",
    "        \"observation\": jnp.zeros((config.num_envs_per_device, 1, *env.observation_shape(env_params))),\n",
    "        \"prev_action\": jnp.zeros((config.num_envs_per_device, 1), dtype=jnp.int32),\n",
    "        \"prev_reward\": jnp.zeros((config.num_envs_per_device, 1)),\n",
    "    }\n",
    "    init_hstate = network.initialize_carry(batch_size=config.num_envs_per_device)\n",
    "\n",
    "    network_params = network.init(_rng, init_obs, init_hstate)\n",
    "    tx = optax.chain(\n",
    "        optax.clip_by_global_norm(config.max_grad_norm),\n",
    "        optax.inject_hyperparams(optax.adam)(learning_rate=linear_schedule, eps=1e-8),  # eps=1e-5\n",
    "    )\n",
    "    train_state = TrainState.create(apply_fn=network.apply, params=network_params, tx=tx)\n",
    "\n",
    "    return rng, env, env_params, benchmark, init_hstate, train_state\n",
    "\n",
    "\n",
    "def make_train(\n",
    "    env: Environment,\n",
    "    env_params: EnvParams,\n",
    "    benchmark: Benchmark,\n",
    "    config: TrainConfig,\n",
    "):\n",
    "    @partial(jax.pmap, axis_name=\"devices\")\n",
    "    def train(\n",
    "        rng: jax.Array,\n",
    "        train_state: TrainState,\n",
    "        init_hstate: jax.Array,\n",
    "    ):\n",
    "        # META TRAIN LOOP\n",
    "        def _meta_step(meta_state, _):\n",
    "            rng, train_state = meta_state\n",
    "\n",
    "            # INIT ENV\n",
    "            rng, _rng1, _rng2 = jax.random.split(rng, num=3)\n",
    "            ruleset_rng = jax.random.split(rng, num=config.num_envs_per_device)\n",
    "            reset_rng = jax.random.split(rng, num=config.num_envs_per_device)\n",
    "\n",
    "            # sample rulesets for this meta update\n",
    "            rulesets = jax.vmap(benchmark.sample_ruleset)(ruleset_rng)\n",
    "            meta_env_params = env_params.replace(ruleset=rulesets)\n",
    "\n",
    "            timestep = jax.vmap(env.reset, in_axes=(0, 0))(meta_env_params, reset_rng)\n",
    "            prev_action = jnp.zeros(config.num_envs_per_device, dtype=jnp.int32)\n",
    "            prev_reward = jnp.zeros(config.num_envs_per_device)\n",
    "\n",
    "            # INNER TRAIN LOOP\n",
    "            def _update_step(runner_state, _):\n",
    "                # COLLECT TRAJECTORIES\n",
    "                def _env_step(runner_state, _):\n",
    "                    rng, train_state, prev_timestep, prev_action, prev_reward, prev_hstate = runner_state\n",
    "\n",
    "                    # SELECT ACTION\n",
    "                    rng, _rng = jax.random.split(rng)\n",
    "                    dist, value, hstate = train_state.apply_fn(\n",
    "                        train_state.params,\n",
    "                        {\n",
    "                            # [batch_size, seq_len=1, ...]\n",
    "                            \"observation\": prev_timestep.observation[:, None],\n",
    "                            \"prev_action\": prev_action[:, None],\n",
    "                            \"prev_reward\": prev_reward[:, None],\n",
    "                        },\n",
    "                        prev_hstate,\n",
    "                    )\n",
    "                    action, log_prob = dist.sample_and_log_prob(seed=_rng)\n",
    "                    # squeeze seq_len where possible\n",
    "                    action, value, log_prob = action.squeeze(1), value.squeeze(1), log_prob.squeeze(1)\n",
    "\n",
    "                    # STEP ENV\n",
    "                    timestep = jax.vmap(env.step, in_axes=0)(meta_env_params, prev_timestep, action)\n",
    "                    transition = Transition(\n",
    "                        # ATTENTION: done is always false, as we optimize for entire meta-rollout\n",
    "                        done=jnp.zeros_like(timestep.last()),\n",
    "                        action=action,\n",
    "                        value=value,\n",
    "                        reward=timestep.reward,\n",
    "                        log_prob=log_prob,\n",
    "                        obs=prev_timestep.observation,\n",
    "                        prev_action=prev_action,\n",
    "                        prev_reward=prev_reward,\n",
    "                    )\n",
    "                    runner_state = (rng, train_state, timestep, action, timestep.reward, hstate)\n",
    "                    return runner_state, transition\n",
    "\n",
    "                initial_hstate = runner_state[-1]\n",
    "                # transitions: [seq_len, batch_size, ...]\n",
    "                runner_state, transitions = jax.lax.scan(_env_step, runner_state, None, config.num_steps_per_update)\n",
    "\n",
    "                # CALCULATE ADVANTAGE\n",
    "                rng, train_state, timestep, prev_action, prev_reward, hstate = runner_state\n",
    "                # calculate value of the last step for bootstrapping\n",
    "                _, last_val, _ = train_state.apply_fn(\n",
    "                    train_state.params,\n",
    "                    {\n",
    "                        \"observation\": timestep.observation[:, None],\n",
    "                        \"prev_action\": prev_action[:, None],\n",
    "                        \"prev_reward\": prev_reward[:, None],\n",
    "                    },\n",
    "                    hstate,\n",
    "                )\n",
    "                advantages, targets = calculate_gae(transitions, last_val.squeeze(1), config.gamma, config.gae_lambda)\n",
    "\n",
    "                # UPDATE NETWORK\n",
    "                def _update_epoch(update_state, _):\n",
    "                    def _update_minbatch(train_state, batch_info):\n",
    "                        init_hstate, transitions, advantages, targets = batch_info\n",
    "                        new_train_state, update_info = ppo_update_networks(\n",
    "                            train_state=train_state,\n",
    "                            transitions=transitions,\n",
    "                            init_hstate=init_hstate.squeeze(1),\n",
    "                            advantages=advantages,\n",
    "                            targets=targets,\n",
    "                            clip_eps=config.clip_eps,\n",
    "                            vf_coef=config.vf_coef,\n",
    "                            ent_coef=config.ent_coef,\n",
    "                        )\n",
    "                        return new_train_state, update_info\n",
    "\n",
    "                    rng, train_state, init_hstate, transitions, advantages, targets = update_state\n",
    "\n",
    "                    # MINIBATCHES PREPARATION\n",
    "                    rng, _rng = jax.random.split(rng)\n",
    "                    permutation = jax.random.permutation(_rng, config.num_envs_per_device)\n",
    "                    # [seq_len, batch_size, ...]\n",
    "                    batch = (init_hstate, transitions, advantages, targets)\n",
    "                    # [batch_size, seq_len, ...], as our model assumes\n",
    "                    batch = jtu.tree_map(lambda x: x.swapaxes(0, 1), batch)\n",
    "\n",
    "                    shuffled_batch = jtu.tree_map(lambda x: jnp.take(x, permutation, axis=0), batch)\n",
    "                    # [num_minibatches, minibatch_size, ...]\n",
    "                    minibatches = jtu.tree_map(\n",
    "                        lambda x: jnp.reshape(x, (config.num_minibatches, -1) + x.shape[1:]), shuffled_batch\n",
    "                    )\n",
    "                    train_state, update_info = jax.lax.scan(_update_minbatch, train_state, minibatches)\n",
    "\n",
    "                    update_state = (rng, train_state, init_hstate, transitions, advantages, targets)\n",
    "                    return update_state, update_info\n",
    "\n",
    "                # hstate shape: [seq_len=None, batch_size, num_layers, hidden_dim]\n",
    "                update_state = (rng, train_state, initial_hstate[None, :], transitions, advantages, targets)\n",
    "                update_state, loss_info = jax.lax.scan(_update_epoch, update_state, None, config.update_epochs)\n",
    "                # WARN: do not forget to get updated params\n",
    "                rng, train_state = update_state[:2]\n",
    "\n",
    "                # averaging over minibatches then over epochs\n",
    "                loss_info = jtu.tree_map(lambda x: x.mean(-1).mean(-1), loss_info)\n",
    "                runner_state = (rng, train_state, timestep, prev_action, prev_reward, hstate)\n",
    "                return runner_state, loss_info\n",
    "\n",
    "            # on each meta-update we reset hidden to init_hstate\n",
    "            runner_state = (rng, train_state, timestep, prev_action, prev_reward, init_hstate)\n",
    "            runner_state, loss_info = jax.lax.scan(_update_step, runner_state, None, config.num_inner_updates)\n",
    "            # WARN: do not forget to get updated params\n",
    "            rng, train_state = runner_state[:2]\n",
    "\n",
    "            # EVALUATE AGENT\n",
    "            eval_ruleset_rng, eval_reset_rng = jax.random.split(jax.random.key(config.eval_seed))\n",
    "            eval_ruleset_rng = jax.random.split(eval_ruleset_rng, num=config.eval_num_envs_per_device)\n",
    "            eval_reset_rng = jax.random.split(eval_reset_rng, num=config.eval_num_envs_per_device)\n",
    "\n",
    "            eval_ruleset = jax.vmap(benchmark.sample_ruleset)(eval_ruleset_rng)\n",
    "            eval_env_params = env_params.replace(ruleset=eval_ruleset)\n",
    "\n",
    "            eval_stats = jax.vmap(rollout, in_axes=(0, None, 0, None, None, None))(\n",
    "                eval_reset_rng,\n",
    "                env,\n",
    "                eval_env_params,\n",
    "                train_state,\n",
    "                # TODO: make this a static method?\n",
    "                jnp.zeros((1, config.rnn_num_layers, config.rnn_hidden_dim)),\n",
    "                config.eval_num_episodes,\n",
    "            )\n",
    "            eval_stats = jax.lax.pmean(eval_stats, axis_name=\"devices\")\n",
    "\n",
    "            # averaging over inner updates, adding evaluation metrics\n",
    "            loss_info = jtu.tree_map(lambda x: x.mean(-1), loss_info)\n",
    "            loss_info.update(\n",
    "                {\n",
    "                    \"eval/returns_mean\": eval_stats.reward.mean(0),\n",
    "                    \"eval/returns_median\": jnp.median(eval_stats.reward),\n",
    "                    \"eval/lengths\": eval_stats.length.mean(0),\n",
    "                    \"eval/lengths_20percentile\": jnp.percentile(eval_stats.length, q=20),\n",
    "                    \"eval/returns_20percentile\": jnp.percentile(eval_stats.reward, q=20),\n",
    "                    \"lr\": train_state.opt_state[-1].hyperparams[\"learning_rate\"],\n",
    "                }\n",
    "            )\n",
    "            meta_state = (rng, train_state)\n",
    "            return meta_state, loss_info\n",
    "\n",
    "        meta_state = (rng, train_state)\n",
    "        meta_state, loss_info = jax.lax.scan(_meta_step, meta_state, None, config.num_meta_updates)\n",
    "        return {\"state\": meta_state[-1], \"loss_info\": loss_info}\n",
    "\n",
    "    return train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb3b5d4c-bcba-4c44-b0f5-0451189dd721",
   "metadata": {},
   "outputs": [],
   "source": [
    "# should take ~8min on colab gpu. It will suboptimal due to the model size and default hyperparams!\n",
    "config = TrainConfig(env_id=\"XLand-MiniGrid-R4-9x9\", benchmark_id=\"trivial-1m\", total_timesteps=100_000_000)\n",
    "\n",
    "rng, env, env_params, benchmark, init_hstate, train_state = make_states(config)\n",
    "# replicating args across devices\n",
    "rng = jax.random.split(rng, num=jax.local_device_count())\n",
    "train_state = replicate(train_state, jax.local_devices())\n",
    "init_hstate = replicate(init_hstate, jax.local_devices())\n",
    "\n",
    "print(\"Compiling...\")\n",
    "t = time.time()\n",
    "train_fn = make_train(env, env_params, benchmark, config)\n",
    "train_fn = train_fn.lower(rng, train_state, init_hstate).compile()\n",
    "elapsed_time = time.time() - t\n",
    "print(f\"Done in {elapsed_time:.2f}s.\")\n",
    "\n",
    "print(\"Training...\")\n",
    "t = time.time()\n",
    "train_info = jax.block_until_ready(train_fn(rng, train_state, init_hstate))\n",
    "elapsed_time = time.time() - t\n",
    "print(f\"Done in {elapsed_time / 60:.2f}min\")\n",
    "\n",
    "# unreplicating from multiple devices\n",
    "train_info = unreplicate(train_info)\n",
    "\n",
    "print(\"Final return: \", float(train_info[\"loss_info\"][\"eval/returns_mean\"][-1]))\n",
    "plt.plot(jnp.arange(config.num_meta_updates), train_info[\"loss_info\"][\"eval/returns_mean\"]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c72527b9-e6f1-4f7e-860b-6e0958c495c7",
   "metadata": {},
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1798edfa-bdeb-4f41-8829-7023215f178e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xminigrid.rendering.text_render import print_ruleset\n",
    "\n",
    "META_EPISODES = 10\n",
    "\n",
    "env, env_params = xminigrid.make(config.env_id)\n",
    "env = GymAutoResetWrapper(env)\n",
    "\n",
    "# enabling image observations if needed\n",
    "if config.img_obs:\n",
    "    from xminigrid.experimental.img_obs import RGBImgObservationWrapper\n",
    "\n",
    "    env = RGBImgObservationWrapper(env)\n",
    "    \n",
    "\n",
    "ruleset = xminigrid.load_benchmark(config.benchmark_id).get_ruleset(ruleset_id=0)\n",
    "env_params = env_params.replace(ruleset=ruleset)\n",
    "\n",
    "# you can use train_state from the final state also\n",
    "# we just demo here how to do it if you loaded params from the checkpoint\n",
    "params = train_info[\"state\"].params\n",
    "model = ActorCriticRNN(\n",
    "    num_actions=env.num_actions(env_params),\n",
    "    action_emb_dim=config.action_emb_dim,\n",
    "    rnn_hidden_dim=config.rnn_hidden_dim,\n",
    "    rnn_num_layers=config.rnn_num_layers,\n",
    "    head_hidden_dim=config.head_hidden_dim,\n",
    "    img_obs=config.img_obs,\n",
    ")\n",
    "\n",
    "# jitting all functions\n",
    "apply_fn, reset_fn, step_fn = jax.jit(model.apply), jax.jit(env.reset), jax.jit(env.step)\n",
    "\n",
    "# for logging\n",
    "total_reward, num_episodes = 0, 0\n",
    "rendered_imgs = []\n",
    "\n",
    "rng = jax.random.key(1)\n",
    "rng, _rng = jax.random.split(rng)\n",
    "\n",
    "# initial inputs\n",
    "hidden = model.initialize_carry(1)\n",
    "prev_reward = jnp.asarray(0)\n",
    "prev_action = jnp.asarray(0)\n",
    "\n",
    "timestep = reset_fn(env_params, _rng)\n",
    "rendered_imgs.append(env.render(env_params, timestep))\n",
    "\n",
    "while num_episodes < META_EPISODES:\n",
    "    rng, _rng = jax.random.split(rng)\n",
    "    dist, _, hidden = apply_fn(\n",
    "        params,\n",
    "        {\n",
    "            \"observation\": timestep.observation[None, None, ...],\n",
    "            \"prev_action\": prev_action[None, None, ...],\n",
    "            \"prev_reward\": prev_reward[None, None, ...],\n",
    "        },\n",
    "        hidden,\n",
    "    )\n",
    "    action = dist.sample(seed=_rng).squeeze()\n",
    "\n",
    "    timestep = step_fn(env_params, timestep, action)\n",
    "    prev_action = action\n",
    "    prev_reward = timestep.reward\n",
    "\n",
    "    total_reward += timestep.reward.item()\n",
    "    num_episodes += int(timestep.last().item())\n",
    "    rendered_imgs.append(env.render(env_params, timestep))\n",
    "\n",
    "print(\"Reward:\", total_reward)\n",
    "print(\"Ruleset:\")\n",
    "print_ruleset(ruleset)\n",
    "imageio.mimsave(\"eval_rollout.mp4\", rendered_imgs, fps=16, format=\"mp4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62eea682-badc-499c-afcb-4d87607bd409",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Video\n",
    "\n",
    "Video(\"eval_rollout.mp4\", embed=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}