memory_sequence.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import flax.linen as nn\n",
    "import numpy as np\n",
    "import optax\n",
    "import time\n",
    "from flax.linen.initializers import constant, orthogonal\n",
    "from typing import Sequence, NamedTuple, Any, Dict, Optional, Union, Tuple\n",
    "from flax.training.train_state import TrainState\n",
    "import distrax\n",
    "import gymnax\n",
    "import functools\n",
    "from functools import partial\n",
    "from gymnax.environments import environment, spaces\n",
    "import chex\n",
    "import wandb\n",
    "from flax import struct\n",
    "\n",
    "\n",
    "class GymnaxWrapper(object):\n",
    "    \"\"\"Base class for Gymnax wrappers.\"\"\"\n",
    "\n",
    "    def __init__(self, env):\n",
    "        self._env = env\n",
    "    \n",
    "    # provide proxy access to regular attributes of wrapped object\n",
    "    def __getattr__(self, name):\n",
    "        return getattr(self._env, name)\n",
    "\n",
    "@struct.dataclass\n",
    "class LogEnvState:\n",
    "    env_state: environment.EnvState\n",
    "    episode_returns: float\n",
    "    episode_lengths: int\n",
    "    returned_episode_returns: float\n",
    "    returned_episode_lengths: int\n",
    "    timestep: int\n",
    "\n",
    "class LogWrapper(GymnaxWrapper):\n",
    "    \"\"\"Log the episode returns and lengths.\"\"\"\n",
    "\n",
    "    def __init__(self, env: environment.Environment):\n",
    "        super().__init__(env)\n",
    "\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def reset(\n",
    "        self, key: chex.PRNGKey, params: Optional[environment.EnvParams] = None\n",
    "    ) -> Tuple[chex.Array, environment.EnvState]:\n",
    "        obs, env_state = self._env.reset(key, params)\n",
    "        state = LogEnvState(env_state, 0, 0, 0, 0, 0)\n",
    "        return obs, state\n",
    "\n",
    "    @partial(jax.jit, static_argnums=(0,))\n",
    "    def step(\n",
    "        self,\n",
    "        key: chex.PRNGKey,\n",
    "        state: environment.EnvState,\n",
    "        action: Union[int, float],\n",
    "        params: Optional[environment.EnvParams] = None,\n",
    "    ) -> Tuple[chex.Array, environment.EnvState, float, bool, dict]:\n",
    "        obs, env_state, reward, done, info = self._env.step(key, state.env_state, action, params)\n",
    "        new_episode_return = state.episode_returns + reward\n",
    "        new_episode_length = state.episode_lengths + 1\n",
    "        state = LogEnvState(\n",
    "            env_state = env_state,\n",
    "            episode_returns = new_episode_return * (1 - done),\n",
    "            episode_lengths = new_episode_length * (1 - done),\n",
    "            returned_episode_returns = state.returned_episode_returns * (1 - done) + new_episode_return * done,\n",
    "            returned_episode_lengths = state.returned_episode_lengths * (1 - done) + new_episode_length * done,\n",
    "            timestep = state.timestep + 1,\n",
    "        )\n",
    "        info[\"returned_episode\"] = done\n",
    "        info[\"return_info\"] = jnp.stack([state.timestep, state.returned_episode_returns])\n",
    "        return obs, state, reward, done, info\n",
    "\n",
    "class ScannedRNN(nn.Module):\n",
    "\n",
    "  @functools.partial(\n",
    "    nn.scan,\n",
    "    variable_broadcast='params',\n",
    "    in_axes=0,\n",
    "    out_axes=0,\n",
    "    split_rngs={'params': False})\n",
    "  @nn.compact\n",
    "  def __call__(self, carry, x):\n",
    "    \"\"\"Applies the module.\"\"\"\n",
    "    features = carry[0].shape[-1]\n",
    "    rnn_state = carry\n",
    "    ins, resets = x\n",
    "    rnn_state = jnp.where(resets[:, np.newaxis], self.initialize_carry(ins.shape[0], ins.shape[1]), rnn_state)\n",
    "    new_rnn_state, y = nn.GRUCell(features)(rnn_state, ins)\n",
    "    return new_rnn_state, y\n",
    "\n",
    "  @staticmethod\n",
    "  def initialize_carry(batch_size, hidden_size):\n",
    "    return nn.GRUCell(hidden_size, parent=None).initialize_carry(\n",
    "        jax.random.PRNGKey(0), (batch_size, hidden_size))\n",
    "\n",
    "class ActorCriticRNN(nn.Module):\n",
    "    action_dim: Sequence[int]\n",
    "    config: Dict\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, hidden, x):\n",
    "        obs, dones = x\n",
    "        embedding = nn.Dense(128, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(obs)\n",
    "        embedding = nn.leaky_relu(embedding)\n",
    "        embedding = nn.Dense(256, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(embedding)\n",
    "        embedding = nn.leaky_relu(embedding)\n",
    "\n",
    "        rnn_in = (embedding, dones)\n",
    "        hidden, embedding = ScannedRNN()(hidden, rnn_in)\n",
    "\n",
    "        actor_mean = nn.Dense(128, kernel_init=orthogonal(2), bias_init=constant(0.0))(embedding)\n",
    "        actor_mean = nn.leaky_relu(actor_mean)\n",
    "        actor_mean = nn.Dense(128, kernel_init=orthogonal(2), bias_init=constant(0.0))(actor_mean)\n",
    "        actor_mean = nn.leaky_relu(actor_mean)\n",
    "        actor_mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(actor_mean)\n",
    "        if self.config[\"CONTINUOUS\"]:\n",
    "            actor_logtstd = self.param('log_std', nn.initializers.zeros, (self.action_dim,))\n",
    "            pi = distrax.MultivariateNormalDiag(actor_mean, jnp.exp(actor_logtstd))\n",
    "        else:\n",
    "            pi = distrax.Categorical(logits=actor_mean)\n",
    "\n",
    "        critic = nn.Dense(128, kernel_init=orthogonal(2), bias_init=constant(0.0))(embedding)\n",
    "        critic = nn.leaky_relu(critic)\n",
    "        critic = nn.Dense(128, kernel_init=orthogonal(2), bias_init=constant(0.0))(critic)\n",
    "        critic = nn.leaky_relu(critic)\n",
    "        critic = nn.Dense(1, kernel_init=orthogonal(1.0), bias_init=constant(0.0))(critic)\n",
    "\n",
    "        return hidden, pi, jnp.squeeze(critic, axis=-1)\n",
    "\n",
    "    @staticmethod\n",
    "    def initialize_carry(batch_size, hidden_size):\n",
    "        return nn.GRUCell(hidden_size, parent=None).initialize_carry(\n",
    "            jax.random.PRNGKey(0), (batch_size, hidden_size))\n",
    "\n",
    "class Transition(NamedTuple):\n",
    "    done: jnp.ndarray\n",
    "    action: jnp.ndarray\n",
    "    value: jnp.ndarray\n",
    "    reward: jnp.ndarray\n",
    "    log_prob: jnp.ndarray\n",
    "    obs: jnp.ndarray\n",
    "    info: jnp.ndarray\n",
    "\n",
    "def make_train(config):\n",
    "    config[\"NUM_UPDATES\"] = (\n",
    "        config[\"TOTAL_TIMESTEPS\"] // config[\"NUM_STEPS\"] // config[\"NUM_ENVS\"]\n",
    "    )\n",
    "    config[\"MINIBATCH_SIZE\"] = (\n",
    "        config[\"NUM_ENVS\"] * config[\"NUM_STEPS\"] // config[\"NUM_MINIBATCHES\"]\n",
    "    )\n",
    "    if \"ENV_NAME\" not in config:\n",
    "        env, env_params = config[\"ENV\"], config[\"ENV_PARAMS\"]\n",
    "        env = LogWrapper(env)\n",
    "    else:\n",
    "        env, env_params = gymnax.make(config[\"ENV_NAME\"])\n",
    "        env = LogWrapper(env)\n",
    "\n",
    "    if \"MOD_ENV_PARAMS\" in config:\n",
    "        env_params = env_params.replace(**config[\"MOD_ENV_PARAMS\"])\n",
    "\n",
    "    config[\"CONTINUOUS\"] = type(env.action_space(env_params)) == spaces.Box \n",
    "\n",
    "    def linear_schedule(count):\n",
    "        frac = 1.0 - (count // (config[\"NUM_MINIBATCHES\"] * config[\"UPDATE_EPOCHS\"])) / config[\"NUM_UPDATES\"]\n",
    "        return config[\"LR\"] * frac\n",
    "\n",
    "    def train(rng):\n",
    "\n",
    "        # INIT NETWORK\n",
    "        if config[\"CONTINUOUS\"]:\n",
    "            network = ActorCriticRNN(env.action_space(env_params).shape[0], config=config)\n",
    "        else:\n",
    "            network = ActorCriticRNN(env.action_space(env_params).n, config=config)\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "        init_x = (jnp.zeros((1, config[\"NUM_ENVS\"], *env.observation_space(env_params).shape)), jnp.zeros((1, config[\"NUM_ENVS\"])))\n",
    "        init_hstate = ScannedRNN.initialize_carry(config[\"NUM_ENVS\"], 256)\n",
    "        network_params = network.init(_rng, init_hstate, init_x)\n",
    "        if config[\"ANNEAL_LR\"]:\n",
    "            tx = optax.chain(\n",
    "                optax.clip_by_global_norm(config[\"MAX_GRAD_NORM\"]),\n",
    "                optax.adam(learning_rate=linear_schedule, eps=1e-5),\n",
    "            )\n",
    "        else:\n",
    "            tx = optax.chain(optax.clip_by_global_norm(config[\"MAX_GRAD_NORM\"]), optax.adam(config[\"LR\"], eps=1e-5))\n",
    "        train_state = TrainState.create(\n",
    "            apply_fn=network.apply,\n",
    "            params=network_params,\n",
    "            tx=tx,\n",
    "        )\n",
    "\n",
    "        # INIT ENV\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "        reset_rng = jax.random.split(_rng, config[\"NUM_ENVS\"])\n",
    "        obsv, env_state = jax.vmap(env.reset, in_axes=(0, None))(reset_rng, env_params)\n",
    "        init_hstate = ScannedRNN.initialize_carry(config[\"NUM_ENVS\"], 256)\n",
    "\n",
    "\n",
    "        # TRAIN LOOP\n",
    "        def _update_step(runner_state, unused):\n",
    "\n",
    "            # COLLECT TRAJECTORIES\n",
    "            def _env_step(runner_state, unused):\n",
    "                train_state, env_state, env_params, last_obs, last_done, hstate, rng = runner_state\n",
    "                if config[\"UPDATE_VIS_PROB\"]:\n",
    "                    env_params = env_params.replace(\n",
    "                        vis_prob=env_params.vis_prob - 1.0 / (config[\"TOTAL_TIMESTEPS\"] / config[\"NUM_ENVS\"])\n",
    "                    )\n",
    "                rng, _rng = jax.random.split(rng)\n",
    "\n",
    "                # SELECT ACTION\n",
    "                ac_in = (last_obs[np.newaxis, :], last_done[np.newaxis, :])\n",
    "                hstate, pi, value = network.apply(train_state.params, hstate, ac_in)\n",
    "                action = pi.sample(seed=_rng)\n",
    "                log_prob = pi.log_prob(action)\n",
    "                value, action, log_prob = value.squeeze(0), action.squeeze(0), log_prob.squeeze(0)\n",
    "\n",
    "                # STEP ENV\n",
    "                rng, _rng = jax.random.split(rng)\n",
    "                rng_step = jax.random.split(_rng, config[\"NUM_ENVS\"])\n",
    "                obsv, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0,0,0,None))(\n",
    "                    rng_step, env_state, action, env_params\n",
    "                )\n",
    "                transition = Transition(last_done, action, value, reward, log_prob, last_obs, info)\n",
    "                runner_state = (train_state, env_state, env_params, obsv, done, hstate, rng)\n",
    "                return runner_state, transition\n",
    "\n",
    "            initial_hstate = runner_state[-2]\n",
    "            runner_state, traj_batch = jax.lax.scan(_env_step, runner_state, None, config[\"NUM_STEPS\"])\n",
    "\n",
    "            # CALCULATE ADVANTAGE\n",
    "            train_state, env_state, env_params, last_obs, last_done, hstate, rng = runner_state\n",
    "            ac_in = (last_obs[np.newaxis, :], last_done[np.newaxis, :])\n",
    "            _, _, last_val = network.apply(train_state.params, hstate, ac_in)\n",
    "            last_val = last_val.squeeze(0)\n",
    "            def _calculate_gae(traj_batch, last_val, last_done):\n",
    "                def _get_advantages(carry, transition):\n",
    "                    gae, next_value, next_done = carry\n",
    "                    done, value, reward = transition.done, transition.value, transition.reward \n",
    "                    delta = reward + config[\"GAMMA\"] * next_value * (1 - next_done) - value\n",
    "                    gae = delta + config[\"GAMMA\"] * config[\"GAE_LAMBDA\"] * (1 - next_done) * gae\n",
    "                    return (gae, value, done), gae\n",
    "                _, advantages = jax.lax.scan(_get_advantages, (jnp.zeros_like(last_val), last_val, last_done), traj_batch, reverse=True, unroll=16)\n",
    "                return advantages, advantages + traj_batch.value\n",
    "            advantages, targets = _calculate_gae(traj_batch, last_val, last_done)\n",
    "\n",
    "            # UPDATE NETWORK\n",
    "            def _update_epoch(update_state, unused):\n",
    "                def _update_minbatch(train_state, batch_info):\n",
    "                    init_hstate, traj_batch,  advantages, targets = batch_info\n",
    "                    def _loss_fn(params, init_hstate, traj_batch, gae, targets):\n",
    "                        # RERUN NETWORK\n",
    "                        _, pi, value = network.apply(params, init_hstate[0], (traj_batch.obs, traj_batch.done))\n",
    "                        log_prob = pi.log_prob(traj_batch.action)\n",
    "\n",
    "                        # CALCULATE VALUE LOSS\n",
    "                        value_pred_clipped = traj_batch.value + (value - traj_batch.value).clip(-config[\"CLIP_EPS\"], config[\"CLIP_EPS\"])\n",
    "                        value_losses = jnp.square(value - targets)\n",
    "                        value_losses_clipped = jnp.square(value_pred_clipped - targets)\n",
    "                        value_loss = 0.5 * jnp.maximum(value_losses, value_losses_clipped).mean()\n",
    "\n",
    "                        # CALCULATE ACTOR LOSS\n",
    "                        ratio = jnp.exp(log_prob - traj_batch.log_prob)\n",
    "                        gae = (gae - gae.mean()) / (gae.std() + 1e-8)\n",
    "                        loss_actor1 = ratio * gae\n",
    "                        loss_actor2 = jnp.clip(ratio, 1.0 - config[\"CLIP_EPS\"], 1.0 + config[\"CLIP_EPS\"]) * gae\n",
    "                        loss_actor = -jnp.minimum(loss_actor1, loss_actor2)\n",
    "                        loss_actor = loss_actor.mean()\n",
    "                        entropy = pi.entropy().mean()\n",
    "\n",
    "                        total_loss = loss_actor + config[\"VF_COEF\"] * value_loss - config[\"ENT_COEF\"] * entropy\n",
    "                        return total_loss, (value_loss, loss_actor, entropy)\n",
    "\n",
    "                    grad_fn = jax.value_and_grad(_loss_fn, has_aux=True)\n",
    "                    total_loss, grads = grad_fn(train_state.params, init_hstate, traj_batch, advantages, targets)\n",
    "                    train_state = train_state.apply_gradients(grads=grads)\n",
    "                    return train_state, total_loss\n",
    "\n",
    "                train_state, init_hstate, traj_batch, advantages, targets, rng = update_state\n",
    "\n",
    "                rng, _rng = jax.random.split(rng)\n",
    "                permutation = jax.random.permutation(_rng, config[\"NUM_ENVS\"])\n",
    "                batch = (init_hstate, traj_batch, advantages, targets)\n",
    "\n",
    "                shuffled_batch = jax.tree_util.tree_map(\n",
    "                    lambda x: jnp.take(x, permutation, axis=1), batch\n",
    "                )\n",
    "\n",
    "                minibatches = jax.tree_util.tree_map(\n",
    "                    lambda x: jnp.swapaxes(jnp.reshape(\n",
    "                        x, [x.shape[0], config[\"NUM_MINIBATCHES\"], -1] + list(x.shape[2:])\n",
    "                    ), 1, 0),\n",
    "                    shuffled_batch,\n",
    "                )\n",
    "\n",
    "                train_state, total_loss = jax.lax.scan(_update_minbatch, train_state, minibatches)\n",
    "                update_state = (train_state, init_hstate, traj_batch, advantages, targets, rng)\n",
    "                return update_state, total_loss\n",
    "\n",
    "            init_hstate = initial_hstate[None,:] # TBH\n",
    "            update_state = (train_state, init_hstate, traj_batch, advantages, targets, rng)\n",
    "            update_state, loss_info = jax.lax.scan(_update_epoch, update_state, None, config[\"UPDATE_EPOCHS\"])\n",
    "            train_state = update_state[0]\n",
    "            metric = traj_batch.info\n",
    "            rng = update_state[-1]\n",
    "\n",
    "            if config[\"DEBUG\"]:\n",
    "                metric = (traj_batch.info[\"return_info\"][...,1]*traj_batch.info[\"returned_episode\"]).sum() / traj_batch.info[\"returned_episode\"].sum()\n",
    "                if config.get(\"LOG\"):\n",
    "                    def callback(metric):\n",
    "                        print(metric)\n",
    "                        wandb.log({\"metric\": metric})\n",
    "                else:\n",
    "                    def callback(metric):\n",
    "                        print(metric)\n",
    "                jax.debug.callback(callback, metric)\n",
    "            else:\n",
    "                metric = (traj_batch.info[\"return_info\"][...,1]*traj_batch.info[\"returned_episode\"]).sum() / traj_batch.info[\"returned_episode\"].sum()\n",
    "\n",
    "            runner_state = (train_state, env_state, env_params, last_obs, last_done, hstate, rng)\n",
    "            return runner_state, metric\n",
    "\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "        runner_state = (train_state, env_state, env_params, obsv, jnp.zeros((config[\"NUM_ENVS\"]), dtype=bool), init_hstate, _rng) \n",
    "        runner_state, metric = jax.lax.scan(_update_step, runner_state, None, config[\"NUM_UPDATES\"])\n",
    "        return runner_state, metric\n",
    "    \n",
    "    return train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IN-WEIGHTS MEMORY ACCUMULATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import lax\n",
    "from gymnax.environments import environment, spaces\n",
    "from typing import Tuple, Optional\n",
    "import chex\n",
    "from flax import struct\n",
    "\n",
    "@struct.dataclass\n",
    "class EnvState:\n",
    "    timestep: int\n",
    "    seq: jnp.ndarray\n",
    "    other_state: jnp.ndarray\n",
    "\n",
    "@struct.dataclass\n",
    "class EnvParams:\n",
    "    other_params: Optional[jnp.ndarray] = None\n",
    "    vis_prob: float = 1.0\n",
    "\n",
    "class MemSeq(environment.Environment):\n",
    "    def __init__(self, num_types=3, max_len=24, other_network=None, num_agents=2.0, init_rng=jax.random.PRNGKey(42), reset_on_mistake=True):\n",
    "        super().__init__()\n",
    "        self.other_network = other_network\n",
    "        self.init_rng = init_rng\n",
    "        self.num_types = num_types\n",
    "        self.max_len = max_len\n",
    "        self.num_agents = num_agents\n",
    "        self.reset_on_mistake = reset_on_mistake\n",
    "\n",
    "    @property\n",
    "    def default_params(self) -> EnvParams:\n",
    "        return EnvParams()\n",
    "    \n",
    "    def step_env(\n",
    "        self, key: chex.PRNGKey, state: EnvState, action: int, params: EnvParams\n",
    "    ) -> Tuple[chex.Array, EnvState, float, bool, dict]:\n",
    "        reward = jnp.where(state.seq[state.timestep] == action, 1.0, -1.0)\n",
    "        terminated = jnp.logical_and(state.seq[state.timestep] != action, self.reset_on_mistake) \n",
    "        # terminated = state.seq[state.timestep] != action\n",
    "        if self.other_network:\n",
    "            key, key_other, key_act = jax.random.split(key, 3)\n",
    "            other_obs = jnp.zeros((self.num_agents * self.num_types+1))\n",
    "            other_obs = other_obs.at[-1].set(1.0)\n",
    "            ac_in = [other_obs[None,None,:], jnp.zeros((1, 1))]\n",
    "            other_state, pi, value = jax.vmap(self.other_network.apply, in_axes=(0,0,None))(params.other_params, state.other_state, ac_in)\n",
    "            other_act = pi.sample(seed=key_act).squeeze(1)\n",
    "        else:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_state = None\n",
    "            other_act = jax.random.randint(key_other, (self.num_agents,), 0, self.num_types)\n",
    "        other_act = other_act.squeeze(-1)\n",
    "        other_act = jax.nn.one_hot(other_act, self.num_types).reshape((-1,))\n",
    "\n",
    "        new_state = EnvState(state.timestep + 1, state.seq, other_state)\n",
    "        terminated = jnp.logical_or(\n",
    "            terminated,\n",
    "            new_state.timestep == self.max_len - 1,\n",
    "        )\n",
    "\n",
    "        key, vis_key = jax.random.split(key)\n",
    "        other_vis = jax.random.uniform(vis_key) < params.vis_prob # TODO: PLR?\n",
    "        other_act = jnp.where(other_vis, other_act, jnp.zeros_like(other_act))\n",
    "        obs = jnp.concatenate([other_act, other_vis[None,]], axis=-1)\n",
    "        return obs, new_state, reward, terminated, {}\n",
    "    \n",
    "    def reset_env(\n",
    "        self, key: chex.PRNGKey, params: EnvParams\n",
    "    ) -> Tuple[chex.Array, EnvState]:\n",
    "        \"\"\"Performs resetting of environment.\"\"\"\n",
    "        if self.init_rng is None:\n",
    "            key, key_seq = jax.random.split(key)\n",
    "            seq = jax.random.randint(key_seq, (self.max_len,), 0, self.num_types)\n",
    "        else:\n",
    "            seq = jax.random.randint(self.init_rng, (self.max_len,), 0, self.num_types)\n",
    "\n",
    "        if self.other_network:\n",
    "            key, key_other, key_act = jax.random.split(key, 3)\n",
    "            other_state = self.other_network.initialize_carry(1, 256) # TODO: LESS HARDCODE\n",
    "            other_state = other_state[None,...].repeat(self.num_agents, axis=0)\n",
    "            other_obs = jnp.zeros((self.num_agents*self.num_types+1))\n",
    "            other_obs = other_obs.at[-1].set(1.0)\n",
    "            ac_in = [other_obs[None,None,:], jnp.ones((1, 1))]\n",
    "            other_state, pi, value = jax.vmap(self.other_network.apply, in_axes=(0,0,None))(params.other_params, other_state, ac_in)\n",
    "            other_act = pi.sample(seed=key_act).squeeze(1)\n",
    "        else:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_state = None\n",
    "            other_act = jax.random.randint(key_other, (self.num_agents,), 0, self.num_types)\n",
    "\n",
    "        other_act = other_act.squeeze(-1)\n",
    "        other_act = jax.nn.one_hot(other_act, self.num_types).reshape((-1,))\n",
    "\n",
    "        state = EnvState(\n",
    "            timestep=0,\n",
    "            seq=seq,\n",
    "            other_state=other_state,\n",
    "        )\n",
    "\n",
    "        key, vis_key = jax.random.split(key)\n",
    "        other_vis = jax.random.uniform(vis_key) < params.vis_prob # TODO: PLR?\n",
    "        other_act = jnp.where(other_vis, other_act, jnp.zeros_like(other_act))\n",
    "        obs = jnp.concatenate([other_act, other_vis[None,]], axis=-1)\n",
    "        return obs, state\n",
    "\n",
    "    def action_space(\n",
    "        self, params: Optional[EnvParams] = None\n",
    "    ) -> spaces.Discrete:\n",
    "        \"\"\"Action space of the environment.\"\"\"\n",
    "        return spaces.Discrete(self.num_types)\n",
    "\n",
    "    def observation_space(self, params: EnvParams) -> spaces.Box:\n",
    "        \"\"\"Observation space of the environment.\"\"\"\n",
    "        return spaces.Box(jnp.zeros((self.num_agents*self.num_types+1,)), jnp.ones((self.num_agents*self.num_types+1,)), (self.num_agents*self.num_types+1,), dtype=jnp.float32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "config = {\n",
    "    \"LR\": 2.5e-4,\n",
    "    \"NUM_ENVS\": 4,\n",
    "    \"NUM_STEPS\": 128,\n",
    "    \"TOTAL_TIMESTEPS\": 1e5,\n",
    "    \"UPDATE_EPOCHS\": 4,\n",
    "    \"NUM_MINIBATCHES\": 4,\n",
    "    \"GAMMA\": 0.99,\n",
    "    \"GAE_LAMBDA\": 0.95,\n",
    "    \"CLIP_EPS\": 0.2,\n",
    "    \"ENT_COEF\": 0.01,\n",
    "    \"VF_COEF\": 0.5,\n",
    "    \"MAX_GRAD_NORM\": 0.5,\n",
    "    \"ANNEAL_LR\": False,\n",
    "    \"DEBUG\": False,\n",
    "    \"CONTINUOUS\": False,\n",
    "    \"UPDATE_VIS_PROB\": True,\n",
    "}\n",
    "\n",
    "popsize = 5\n",
    "num_types = 10\n",
    "rng = jax.random.PRNGKey(64)\n",
    "\n",
    "env = MemSeq(\n",
    "    num_types=num_types,\n",
    "    other_network=ActorCriticRNN(num_types, config=config),\n",
    "    num_agents=popsize,\n",
    ")\n",
    "\n",
    "rng, rng_init = jax.random.split(rng)\n",
    "rng_inits = jax.random.split(rng_init, popsize)\n",
    "init_other_params = jax.vmap(env.other_network.init, in_axes=(0, None, None))(rng_inits, ScannedRNN.initialize_carry(1, 256), (jnp.zeros((1, 1, popsize*num_types + 1)), jnp.ones((1, 1))))\n",
    "config[\"ENV\"] = env\n",
    "env_params = EnvParams(other_params=init_other_params)\n",
    "config[\"ENV_PARAMS\"] = env_params\n",
    "\n",
    "def train_wparams(rng, other_params):\n",
    "    config_copy = config.copy()\n",
    "    config_copy[\"ENV_PARAMS\"] = EnvParams(other_params=other_params)\n",
    "    train = make_train(config_copy)\n",
    "    return train(rng)\n",
    "\n",
    "jit_vmap_train = jax.jit(jax.vmap(train_wparams, in_axes=(0, None)))\n",
    "\n",
    "all_out_metrics = []\n",
    "other_params = init_other_params\n",
    "for i in range(8):\n",
    "    print(\"*\"*32)\n",
    "    print(f\"NEW GENERATION {i}\")\n",
    "    rng, rng_train = jax.random.split(rng)\n",
    "    rngs = jax.random.split(rng_train, popsize)\n",
    "    outs = jit_vmap_train(rngs, other_params)\n",
    "    print(outs[1][...,-1])\n",
    "    all_out_metrics.append(outs[1])\n",
    "    best_sorted = jnp.argsort(outs[1][...,-1])\n",
    "    other_params = jax.tree_util.tree_map(lambda x: x[best_sorted], outs[0][0].params)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "\n",
    "members = all_out_metrics\n",
    "n = len(members)\n",
    "colors = cm.viridis([i/n for i in range(n)])  # Use the viridis colormap\n",
    "\n",
    "best_vals = []\n",
    "for i, (m, color) in enumerate(zip(members, colors)):\n",
    "    best = jnp.argmax(m[...,-1])\n",
    "    best_vals.append(m[best][...,-1])\n",
    "    m = m[best]\n",
    "    plt.plot(m, label=f\"Generation {i}\", color=color, alpha=0.5)\n",
    "\n",
    "plt.legend()\n",
    "plt.ylabel(\"Return\")\n",
    "plt.xlabel(\"Update\")\n",
    "plt.title(\"MemSeq\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IN-CONTEXT MEMORY ACCUMULATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import lax\n",
    "from gymnax.environments import environment, spaces\n",
    "from typing import Tuple, Optional\n",
    "import chex\n",
    "from flax import struct\n",
    "\n",
    "@struct.dataclass\n",
    "class EnvState:\n",
    "    timestep: int\n",
    "    seq: jnp.ndarray\n",
    "    trial_num: int\n",
    "    trial_timestep: int\n",
    "    # OTHER NETWORK IN-CONTEXT\n",
    "    other_state: Optional[jnp.ndarray] = None\n",
    "    other_act: Optional[jnp.ndarray] = None\n",
    "    # NO NETWORK OTHERS\n",
    "    other_probs: Optional[jnp.array] = None\n",
    "\n",
    "@struct.dataclass\n",
    "class EnvParams:\n",
    "    # OTHER NETWORK IN-CONTEXT\n",
    "    other_params: Optional[jnp.ndarray] = None\n",
    "    other_init_state: Optional[jnp.ndarray] = None\n",
    "    other_init_last_obs: Optional[jnp.ndarray] = None\n",
    "    # OTHER MAX\n",
    "    other_max_len: Optional[int] = None\n",
    "\n",
    "class MetaMemSeq(environment.Environment):\n",
    "    def __init__(\n",
    "            self, \n",
    "            num_types=3,\n",
    "            max_trial_len=24,\n",
    "            num_trials=10,\n",
    "            other_network=None,\n",
    "            init_rng=None,\n",
    "            num_agents=1,\n",
    "            sort_best=True,\n",
    "            reset_on_mistake=True,\n",
    "            other_probs=None,\n",
    "        ):\n",
    "        super().__init__()\n",
    "        self.num_types = num_types\n",
    "        self.max_trial_len = max_trial_len\n",
    "        self.num_trials = num_trials\n",
    "        self.other_network = other_network\n",
    "        self.init_rng = init_rng\n",
    "        self.num_agents = num_agents\n",
    "        self.sort_best = sort_best\n",
    "        self.reset_on_mistake = reset_on_mistake\n",
    "        self.other_probs = other_probs\n",
    "\n",
    "    @property\n",
    "    def default_params(self) -> EnvParams:\n",
    "        return EnvParams()\n",
    "\n",
    "    def step_env(\n",
    "        self, key: chex.PRNGKey, state: EnvState, action: int, params: EnvParams\n",
    "    ) -> Tuple[chex.Array, EnvState, float, bool, dict]:\n",
    "        reward = jnp.where(state.seq[state.trial_timestep] == action, 1.0, -1.0)\n",
    "        trial_terminated = jnp.logical_and(state.seq[state.trial_timestep] != action, self.reset_on_mistake)\n",
    "        trial_terminated = jnp.logical_or(\n",
    "            trial_terminated,\n",
    "            state.trial_timestep == self.max_trial_len - 1,\n",
    "        )\n",
    "\n",
    "        next_trial_timestep = jnp.where(trial_terminated, 0, state.trial_timestep + 1)\n",
    "        next_trial_num = jnp.where(trial_terminated, state.trial_num + 1, state.trial_num)\n",
    "        terminated = jnp.logical_and(\n",
    "            trial_terminated,\n",
    "            next_trial_num == self.num_trials,\n",
    "        ) \n",
    "\n",
    "        if self.other_network is None:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_act_rand = jax.random.randint(key_other, (self.num_agents,), 0, self.num_types)\n",
    "            other_act_gt = state.seq[next_trial_timestep][None,...]\n",
    "            key, key_other_gt = jax.random.split(key)\n",
    "            key_other_gts = jax.random.split(key_other_gt, self.num_agents)\n",
    "            other_gt = jax.vmap(jax.random.uniform)(key_other_gts) < state.other_probs\n",
    "            other_act = jnp.where(other_gt, other_act_gt, other_act_rand)\n",
    "            other_state = None\n",
    "        else:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_state = jnp.where(trial_terminated, params.other_init_state, state.other_state)\n",
    "            # other_last_act = jnp.where(trial_terminated, params.other_init_last_act, action)\n",
    "            # other_reward = jnp.where(trial_terminated, params.other_init_last_reward, reward)\n",
    "            other_obs = jnp.concatenate([\n",
    "                jax.nn.one_hot(action, self.num_types),\n",
    "                trial_terminated[None,],\n",
    "                reward[None,],\n",
    "                jnp.zeros(self.num_agents*self.num_types+1,),\n",
    "            ], axis=-1)[None,:]\n",
    "            # print(other_obs.shape)\n",
    "            other_obs = jnp.tile(other_obs, (self.num_agents, 1))\n",
    "            # print(other_obs.shape)\n",
    "            other_obs = jnp.where(trial_terminated, params.other_init_last_obs, other_obs)\n",
    "\n",
    "            # ac_in = [other_obs[None,None,:], jnp.zeros((1, 1))]\n",
    "            ac_in = [other_obs[:,None,None,:], jnp.zeros((self.num_agents, 1, 1))]\n",
    "            other_state, pi, value = jax.vmap(self.other_network.apply, in_axes=(None,0,0))(params.other_params, other_state, ac_in)\n",
    "            other_act = pi.sample(seed=key_other).squeeze(1).squeeze(1)\n",
    "            # other_act = jax.nn.one_hot(other_act, self.num_types)\n",
    "\n",
    "\n",
    "        own_act = jax.nn.one_hot(action, self.num_types)\n",
    "\n",
    "        new_state = EnvState(\n",
    "            state.timestep + 1,\n",
    "            state.seq,\n",
    "            next_trial_num, \n",
    "            next_trial_timestep, \n",
    "            other_state, \n",
    "            other_act,\n",
    "            state.other_probs,\n",
    "        )\n",
    "\n",
    "        key, vis_key = jax.random.split(key)\n",
    "        other_vis = jax.random.uniform(vis_key) < (self.num_trials - state.trial_num - 1) / (self.num_trials-1) # TODO: PLR?\n",
    "        other_act = jax.nn.one_hot(other_act, self.num_types)\n",
    "        other_act = jnp.where(other_vis, other_act, jnp.zeros_like(other_act)).reshape((-1,))\n",
    "        obs = jnp.concatenate([own_act, trial_terminated[None,], reward[None,], other_act, other_vis[None,]], axis=-1)\n",
    "        return obs, new_state, reward, terminated, {}\n",
    "\n",
    "    def reset_env(\n",
    "        self, key: chex.PRNGKey, params: EnvParams\n",
    "    ) -> Tuple[chex.Array, EnvState]:\n",
    "        \"\"\"Performs resetting of environment.\"\"\"\n",
    "        if self.init_rng is None:\n",
    "            key, key_seq = jax.random.split(key)\n",
    "            seq = jax.random.randint(key_seq, (self.max_trial_len,), 0, self.num_types)\n",
    "        else:\n",
    "            seq = jax.random.randint(self.init_rng, (self.max_trial_len,), 0, self.num_types)\n",
    "        \n",
    "        if self.other_network is None:\n",
    "            if self.other_probs is not None:\n",
    "                other_probs = self.other_probs\n",
    "            else:\n",
    "                key, key_prob = jax.random.split(key)\n",
    "                key_probs = jax.random.split(key_prob, self.num_agents)\n",
    "                if params.other_max_len is None:\n",
    "                    print(\"NO MAX LEN\")\n",
    "                    other_probs = jax.vmap(jax.random.uniform)(key_probs) * 3 + 1\n",
    "                else:\n",
    "                    other_probs = jax.vmap(jax.random.uniform)(key_probs) * params.other_max_len + 1\n",
    "                other_probs = 1.0 - 1.0 / other_probs\n",
    "            if self.sort_best:\n",
    "                other_probs = jnp.sort(other_probs)\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_act_rand = jax.random.randint(key_other, (self.num_agents,), 0, self.num_types)\n",
    "            other_act_gt = seq[0][None,...]\n",
    "            key, key_other_gt = jax.random.split(key)\n",
    "            key_other_gts = jax.random.split(key_other_gt, self.num_agents)\n",
    "            other_gt = jax.vmap(jax.random.uniform)(key_other_gts) < other_probs\n",
    "            other_act = jnp.where(other_gt, other_act_gt, other_act_rand)\n",
    "            other_state = None\n",
    "        else:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_obs = params.other_init_last_obs\n",
    "            ac_in = [other_obs[:,None,None,:], jnp.zeros((self.num_agents, 1, 1))]\n",
    "            other_state, pi, value = jax.vmap(self.other_network.apply, in_axes=(None, 0, 0))(params.other_params, params.other_init_state, ac_in)\n",
    "            other_act = pi.sample(seed=key_other).squeeze(1).squeeze(1)\n",
    "            other_probs = None\n",
    "\n",
    "        state = EnvState(\n",
    "            timestep=0,\n",
    "            seq=seq,\n",
    "            trial_num=0,\n",
    "            trial_timestep=0,\n",
    "            other_state=other_state,\n",
    "            other_act=other_act,\n",
    "            other_probs=other_probs,\n",
    "        )\n",
    "\n",
    "        key, vis_key = jax.random.split(key)\n",
    "        other_vis = jax.random.uniform(vis_key) < (self.num_trials - state.trial_num - 1) / (self.num_trials - 1)\n",
    "        other_act = jax.nn.one_hot(other_act, self.num_types)\n",
    "        other_act = jnp.where(other_vis, other_act, jnp.zeros_like(other_act)).reshape((-1,))\n",
    "        obs = jnp.concatenate([jnp.zeros(self.num_types,), jnp.zeros((2,)), other_act, other_vis[None,]], axis=-1)\n",
    "        return obs, state\n",
    "    \n",
    "    def action_space(\n",
    "        self, params: Optional[EnvParams] = None\n",
    "    ) -> spaces.Discrete:\n",
    "        \"\"\"Action space of the environment.\"\"\"\n",
    "        return spaces.Discrete(self.num_types)\n",
    "\n",
    "    def observation_space(self, params: EnvParams) -> spaces.Box:\n",
    "        \"\"\"Observation space of the environment.\"\"\"\n",
    "        return spaces.Box(jnp.zeros((self.num_types*(self.num_agents+1)+3,)), jnp.ones((self.num_types*(self.num_agents+1)+3,)), (self.num_types*(self.num_agents+1)+3,), dtype=jnp.float32)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_types = 3\n",
    "num_trials = 4\n",
    "max_trial_len = 24\n",
    "popsize = 3\n",
    "\n",
    "def train_batch_(rng, max_len):\n",
    "    config = {\n",
    "        \"LR\": 2.5e-4,\n",
    "        \"NUM_ENVS\": 16,\n",
    "        \"NUM_STEPS\": 128,\n",
    "        \"TOTAL_TIMESTEPS\": 8e6,\n",
    "        \"UPDATE_EPOCHS\": 2,\n",
    "        \"NUM_MINIBATCHES\": 8,\n",
    "        \"GAMMA\": 0.99,\n",
    "        \"GAE_LAMBDA\": 0.95,\n",
    "        \"CLIP_EPS\": 0.2,\n",
    "        \"ENT_COEF\": 0.03,\n",
    "        \"VF_COEF\": 0.5,\n",
    "        \"MAX_GRAD_NORM\": 0.5,\n",
    "        \"ANNEAL_LR\": True,\n",
    "        \"DEBUG\": False,\n",
    "        \"CONTINUOUS\": False,\n",
    "        \"UPDATE_VIS_PROB\": False,\n",
    "    }\n",
    "\n",
    "    env = MetaMemSeq(\n",
    "        num_types=num_types,\n",
    "        num_trials=num_trials,\n",
    "        max_trial_len=max_trial_len,\n",
    "        num_agents=popsize,\n",
    "    )\n",
    "\n",
    "    config[\"ENV\"] = env\n",
    "    config[\"ENV_PARAMS\"] = EnvParams(\n",
    "        other_max_len=max_len,\n",
    "    )\n",
    "    train = make_train(config)\n",
    "    out = train(rng)\n",
    "    return out\n",
    "\n",
    "max_lens = jnp.array([0.5])\n",
    "num_seeds = 5\n",
    "rng = jax.random.PRNGKey(64)\n",
    "rngs = jax.random.split(rng, num_seeds)\n",
    "all_outs = jax.jit(jax.vmap(jax.vmap(train_batch_, in_axes=(None, 0)), in_axes=(0, None)))(rngs, max_lens)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "for i in range(len(max_lens)):\n",
    "    plt.plot(all_outs[1][:,i].mean(0), label=f\"Max Len {i+1}\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ICL GENERAITION EVAL\n",
    "config = {\n",
    "    \"LR\": 2.5e-4,\n",
    "    \"NUM_ENVS\": 16,\n",
    "    \"NUM_STEPS\": 128,\n",
    "    \"TOTAL_TIMESTEPS\": 8e6,\n",
    "    \"UPDATE_EPOCHS\": 2,\n",
    "    \"NUM_MINIBATCHES\": 8,\n",
    "    \"GAMMA\": 0.99,\n",
    "    \"GAE_LAMBDA\": 0.95,\n",
    "    \"CLIP_EPS\": 0.2,\n",
    "    \"ENT_COEF\": 0.03,\n",
    "    \"VF_COEF\": 0.5,\n",
    "    \"MAX_GRAD_NORM\": 0.5,\n",
    "    \"ANNEAL_LR\": True,\n",
    "    \"DEBUG\": False,\n",
    "    \"CONTINUOUS\": False,\n",
    "    \"UPDATE_VIS_PROB\": False,\n",
    "}\n",
    "num_types = 3\n",
    "num_trials = 4\n",
    "max_trial_len = 24\n",
    "popsize = 3\n",
    "\n",
    "init_rng = jax.random.PRNGKey(420)\n",
    "init_env = MetaMemSeq(\n",
    "    num_types=num_types,\n",
    "    init_rng=init_rng,\n",
    "    num_trials=num_trials,\n",
    "    max_trial_len=max_trial_len,\n",
    "    num_agents=popsize,\n",
    "    other_probs=jnp.zeros((popsize,)),\n",
    ")\n",
    "\n",
    "network = ActorCriticRNN(num_types, config=config)\n",
    "\n",
    "env = MetaMemSeq(\n",
    "    num_types=num_types,\n",
    "    other_network=network,\n",
    "    init_rng=init_rng,\n",
    "    num_trials=num_trials,\n",
    "    max_trial_len=max_trial_len,\n",
    "    num_agents=popsize,\n",
    ")\n",
    "\n",
    "init_env_params = EnvParams()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_eval_params_and_return(rng, params, env_params):\n",
    "    rng, _rng = jax.random.split(rng)\n",
    "    last_obs, env_state = init_env.reset(_rng, env_params)\n",
    "\n",
    "    init_state = ScannedRNN.initialize_carry(1, 256)\n",
    "\n",
    "    # COLLECT TRAJECTORIES\n",
    "    def _env_step(runner_state, unused):\n",
    "        prev_env_state, env_params, last_obs, last_done, prev_hstate, rng, ever_done, running_r, running_hstate, running_obs = runner_state\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "\n",
    "        # SELECT ACTION\n",
    "        ac_in = (last_obs[None, None, :], last_done[None, None])\n",
    "        hstate, pi, value = network.apply(params, prev_hstate, ac_in)\n",
    "        action = pi.sample(seed=_rng).squeeze(0).squeeze(0)\n",
    "\n",
    "        # STEP ENV\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "        obsv, env_state, reward, done, info = init_env.step(\n",
    "            _rng, prev_env_state, action, env_params\n",
    "        )\n",
    "\n",
    "        # ONLY GET REWARDS FROM FIRST META-EPISODE\n",
    "        temp_r = running_r[prev_env_state.trial_num] + (reward >= 0.0) * (1.0 - ever_done)\n",
    "        running_r = running_r.at[prev_env_state.trial_num].set(temp_r)\n",
    "\n",
    "        # ONLY FINAL HIDDEN STATE AND ACTION\n",
    "        replace_hstate = jnp.logical_and((prev_env_state.trial_num == env.num_trials - 2), ~ever_done)\n",
    "        running_hstate = running_hstate * (1.0 - replace_hstate) + hstate * replace_hstate\n",
    "        running_obs = running_obs * (1.0 - replace_hstate) + obsv * replace_hstate\n",
    "\n",
    "        ever_done = jnp.logical_or(done, ever_done)\n",
    "        # transition = (prev_env_state, action, running_r, reward, hstate)\n",
    "        transition = (env_state, action, reward, replace_hstate, running_r, running_hstate, running_obs)\n",
    "        runner_state = (env_state, env_params, obsv, done, hstate, rng, ever_done, running_r, running_hstate, running_obs)\n",
    "        return runner_state, transition\n",
    "\n",
    "    runner_state = (env_state, env_params, last_obs, False, init_state, rng, False, jnp.zeros((env.num_trials,)), init_state, last_obs)\n",
    "    runner_state, traj_batch = jax.lax.scan(_env_step, runner_state, None, env.max_trial_len*env.num_trials+1)\n",
    "    return traj_batch\n",
    "\n",
    "def eval_params_and_return(rng, params, env_params):\n",
    "    rng, _rng = jax.random.split(rng)\n",
    "    last_obs, env_state = env.reset(_rng, env_params)\n",
    "\n",
    "    init_state = ScannedRNN.initialize_carry(1, 256)\n",
    "\n",
    "    # COLLECT TRAJECTORIES\n",
    "    def _env_step(runner_state, unused):\n",
    "        prev_env_state, env_params, last_obs, last_done, prev_hstate, rng, ever_done, running_r, running_hstate, running_obs = runner_state\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "\n",
    "        # SELECT ACTION\n",
    "        ac_in = (last_obs[None, None, :], last_done[None, None])\n",
    "        hstate, pi, value = network.apply(params, prev_hstate, ac_in)\n",
    "        action = pi.sample(seed=_rng).squeeze(0).squeeze(0)\n",
    "\n",
    "        # STEP ENV\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "        obsv, env_state, reward, done, info = env.step(\n",
    "            _rng, prev_env_state, action, env_params\n",
    "        )\n",
    "\n",
    "        # ONLY GET REWARDS FROM FIRST META-EPISODE\n",
    "        temp_r = running_r[prev_env_state.trial_num] + (reward >= 0.0) * (1.0 - ever_done)\n",
    "        running_r = running_r.at[prev_env_state.trial_num].set(temp_r)\n",
    "\n",
    "        # ONLY FINAL HIDDEN STATE AND ACTION\n",
    "        replace_hstate = jnp.logical_and((prev_env_state.trial_num == env.num_trials - 2), ~ever_done)\n",
    "        running_hstate = running_hstate * (1.0 - replace_hstate) + hstate * replace_hstate\n",
    "        running_obs = running_obs * (1.0 - replace_hstate) + obsv * replace_hstate\n",
    "\n",
    "        ever_done = jnp.logical_or(done, ever_done)\n",
    "        # transition = (prev_env_state, action, running_r, reward, hstate)\n",
    "        transition = (env_state, action, reward, running_r, running_hstate, running_obs)\n",
    "        runner_state = (env_state, env_params, obsv, done, hstate, rng, ever_done, running_r, running_hstate, running_obs)\n",
    "        return runner_state, transition\n",
    "\n",
    "    runner_state = (env_state, env_params, last_obs, False, init_state, rng, False, jnp.zeros((env.num_trials,)), init_state, last_obs)\n",
    "    runner_state, traj_batch = jax.lax.scan(_env_step, runner_state, None, env.max_trial_len*env.num_trials+1)\n",
    "    return traj_batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_all_trial_scores = []\n",
    "for k in range(num_seeds):\n",
    "    all_trial_scores = []\n",
    "\n",
    "    for i in range(len(max_lens)):\n",
    "        print(f\"Seed {k}, Max Len {i+1}\")\n",
    "\n",
    "        rng = jax.random.PRNGKey(12)\n",
    "        test = jax.jit(jax.vmap(init_eval_params_and_return, in_axes=(0, None, None)))\n",
    "        other_params = jax.tree_map(lambda x: x[k,i], all_outs[0][0].params)\n",
    "        infos = test(jax.random.split(rng, popsize), other_params, init_env_params)\n",
    "\n",
    "        num_generations = 8\n",
    "        trial_scores = []\n",
    "        for j in range(num_generations):\n",
    "            last_trial_scores = infos[-3][:,-1,-1]\n",
    "            last_trial_scores_idx = jnp.argsort(last_trial_scores)\n",
    "            saved_scores = infos[-3][:,-1,:][last_trial_scores_idx]\n",
    "            saved_states = infos[-2][:,-1,:][last_trial_scores_idx]\n",
    "            saved_obs = infos[-1][:,-1,:][last_trial_scores_idx]\n",
    "            trial_scores.append(saved_scores.mean(0))\n",
    "            print(saved_scores.mean(0))\n",
    "\n",
    "            env_params = EnvParams(\n",
    "                other_params=other_params,\n",
    "                other_init_state=saved_states,\n",
    "                other_init_last_obs=saved_obs,\n",
    "            )\n",
    "\n",
    "            rng = jax.random.PRNGKey(3)\n",
    "            test = jax.jit(jax.vmap(eval_params_and_return, in_axes=(0, None, None)))\n",
    "            infos = test(jax.random.split(rng, popsize), other_params, env_params)\n",
    "\n",
    "        trial_scores.append(saved_scores.mean(0))\n",
    "        all_trial_scores.append(np.array(trial_scores))\n",
    "    test_all_trial_scores.append(np.array(all_trial_scores))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RL^2 BASELINE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import lax\n",
    "from gymnax.environments import environment, spaces\n",
    "from typing import Tuple, Optional\n",
    "import chex\n",
    "from flax import struct\n",
    "\n",
    "@struct.dataclass\n",
    "class EnvState:\n",
    "    timestep: int\n",
    "    seq: jnp.ndarray\n",
    "    trial_num: int\n",
    "    trial_timestep: int\n",
    "    # OTHER NETWORK IN-CONTEXT\n",
    "    other_state: Optional[jnp.ndarray] = None\n",
    "    other_act: Optional[jnp.ndarray] = None\n",
    "    # NO NETWORK OTHERS\n",
    "    other_probs: Optional[jnp.array] = None\n",
    "\n",
    "@struct.dataclass\n",
    "class EnvParams:\n",
    "    # OTHER NETWORK IN-CONTEXT\n",
    "    other_params: Optional[jnp.ndarray] = None\n",
    "    other_init_state: Optional[jnp.ndarray] = None\n",
    "    other_init_last_obs: Optional[jnp.ndarray] = None\n",
    "    # OTHER MAX\n",
    "    other_max_len: Optional[int] = None\n",
    "\n",
    "class MetaMemSeq(environment.Environment):\n",
    "    def __init__(\n",
    "            self, \n",
    "            num_types=3,\n",
    "            max_trial_len=24,\n",
    "            num_trials=10,\n",
    "            other_network=None,\n",
    "            init_rng=None,\n",
    "            num_agents=1,\n",
    "            sort_best=True,\n",
    "            reset_on_mistake=True,\n",
    "            other_probs=None,\n",
    "        ):\n",
    "        super().__init__()\n",
    "        self.num_types = num_types\n",
    "        self.max_trial_len = max_trial_len\n",
    "        self.num_trials = num_trials\n",
    "        self.other_network = other_network\n",
    "        self.init_rng = init_rng\n",
    "        self.num_agents = num_agents\n",
    "        self.sort_best = sort_best\n",
    "        self.reset_on_mistake = reset_on_mistake\n",
    "        self.other_probs = other_probs\n",
    "\n",
    "    @property\n",
    "    def default_params(self) -> EnvParams:\n",
    "        return EnvParams()\n",
    "\n",
    "    def step_env(\n",
    "        self, key: chex.PRNGKey, state: EnvState, action: int, params: EnvParams\n",
    "    ) -> Tuple[chex.Array, EnvState, float, bool, dict]:\n",
    "        reward = jnp.where(state.seq[state.trial_timestep] == action, 1.0, -1.0)\n",
    "        trial_terminated = jnp.logical_and(state.seq[state.trial_timestep] != action, self.reset_on_mistake)\n",
    "        trial_terminated = jnp.logical_or(\n",
    "            trial_terminated,\n",
    "            state.trial_timestep == self.max_trial_len - 1,\n",
    "        )\n",
    "\n",
    "        next_trial_timestep = jnp.where(trial_terminated, 0, state.trial_timestep + 1)\n",
    "        next_trial_num = jnp.where(trial_terminated, state.trial_num + 1, state.trial_num)\n",
    "        terminated = jnp.logical_and(\n",
    "            trial_terminated,\n",
    "            next_trial_num == self.num_trials,\n",
    "        ) \n",
    "\n",
    "        if self.other_network is None:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_act_rand = jax.random.randint(key_other, (self.num_agents,), 0, self.num_types)\n",
    "            other_act_gt = state.seq[next_trial_timestep][None,...]\n",
    "            key, key_other_gt = jax.random.split(key)\n",
    "            key_other_gts = jax.random.split(key_other_gt, self.num_agents)\n",
    "            other_gt = jax.vmap(jax.random.uniform)(key_other_gts) < state.other_probs\n",
    "            other_act = jnp.where(other_gt, other_act_gt, other_act_rand)\n",
    "            other_state = None\n",
    "        else:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_state = jnp.where(trial_terminated, params.other_init_state, state.other_state)\n",
    "            other_obs = jnp.concatenate([\n",
    "                jax.nn.one_hot(action, self.num_types),\n",
    "                trial_terminated[None,],\n",
    "                reward[None,],\n",
    "                jnp.zeros(self.num_agents*self.num_types+1,),\n",
    "            ], axis=-1)[None,:]\n",
    "            other_obs = jnp.tile(other_obs, (self.num_agents, 1))\n",
    "            other_obs = jnp.where(trial_terminated, params.other_init_last_obs, other_obs)\n",
    "\n",
    "            ac_in = [other_obs[:,None,None,:], jnp.zeros((self.num_agents, 1, 1))]\n",
    "            other_state, pi, value = jax.vmap(self.other_network.apply, in_axes=(None,0,0))(params.other_params, other_state, ac_in)\n",
    "            other_act = pi.sample(seed=key_other).squeeze(1).squeeze(1)\n",
    "\n",
    "\n",
    "        own_act = jax.nn.one_hot(action, self.num_types)\n",
    "\n",
    "        new_state = EnvState(\n",
    "            state.timestep + 1,\n",
    "            state.seq,\n",
    "            next_trial_num, \n",
    "            next_trial_timestep, \n",
    "            other_state, \n",
    "            other_act,\n",
    "            state.other_probs,\n",
    "        )\n",
    "\n",
    "        key, vis_key = jax.random.split(key)\n",
    "        other_vis = jax.random.uniform(vis_key) < (self.num_trials - state.trial_num - 1) / (self.num_trials-1) # TODO: PLR?\n",
    "        other_act = jax.nn.one_hot(other_act, self.num_types)\n",
    "        other_act = jnp.where(other_vis, other_act, jnp.zeros_like(other_act)).reshape((-1,))\n",
    "        obs = jnp.concatenate([own_act, trial_terminated[None,], reward[None,], other_act, other_vis[None,]], axis=-1)\n",
    "        return obs, new_state, reward, terminated, {}\n",
    "\n",
    "    def reset_env(\n",
    "        self, key: chex.PRNGKey, params: EnvParams\n",
    "    ) -> Tuple[chex.Array, EnvState]:\n",
    "        \"\"\"Performs resetting of environment.\"\"\"\n",
    "        if self.init_rng is None:\n",
    "            key, key_seq = jax.random.split(key)\n",
    "            seq = jax.random.randint(key_seq, (self.max_trial_len,), 0, self.num_types)\n",
    "        else:\n",
    "            seq = jax.random.randint(self.init_rng, (self.max_trial_len,), 0, self.num_types)\n",
    "        \n",
    "        if self.other_network is None:\n",
    "            if self.other_probs is not None:\n",
    "                other_probs = self.other_probs\n",
    "            else:\n",
    "                key, key_prob = jax.random.split(key)\n",
    "                key_probs = jax.random.split(key_prob, self.num_agents)\n",
    "                if params.other_max_len is None:\n",
    "                    print(\"NO MAX LEN\")\n",
    "                    other_probs = jax.vmap(jax.random.uniform)(key_probs) * 3 + 1\n",
    "                else:\n",
    "                    other_probs = jax.vmap(jax.random.uniform)(key_probs) * params.other_max_len + 1\n",
    "                other_probs = 1.0 - 1.0 / other_probs\n",
    "            if self.sort_best:\n",
    "                other_probs = jnp.sort(other_probs)\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_act_rand = jax.random.randint(key_other, (self.num_agents,), 0, self.num_types)\n",
    "            other_act_gt = seq[0][None,...]\n",
    "            key, key_other_gt = jax.random.split(key)\n",
    "            key_other_gts = jax.random.split(key_other_gt, self.num_agents)\n",
    "            other_gt = jax.vmap(jax.random.uniform)(key_other_gts) < other_probs\n",
    "            other_act = jnp.where(other_gt, other_act_gt, other_act_rand)\n",
    "            other_state = None\n",
    "        else:\n",
    "            key, key_other = jax.random.split(key)\n",
    "            other_obs = params.other_init_last_obs\n",
    "            ac_in = [other_obs[:,None,None,:], jnp.zeros((self.num_agents, 1, 1))]\n",
    "            other_state, pi, value = jax.vmap(self.other_network.apply, in_axes=(None, 0, 0))(params.other_params, params.other_init_state, ac_in)\n",
    "            other_act = pi.sample(seed=key_other).squeeze(1).squeeze(1)\n",
    "            other_probs = None\n",
    "\n",
    "        state = EnvState(\n",
    "            timestep=0,\n",
    "            seq=seq,\n",
    "            trial_num=0,\n",
    "            trial_timestep=0,\n",
    "            other_state=other_state,\n",
    "            other_act=other_act,\n",
    "            other_probs=other_probs,\n",
    "        )\n",
    "\n",
    "        key, vis_key = jax.random.split(key)\n",
    "        other_vis = jax.random.uniform(vis_key) < (self.num_trials - state.trial_num - 1) / (self.num_trials - 1)\n",
    "        other_act = jax.nn.one_hot(other_act, self.num_types)\n",
    "        other_act = jnp.where(other_vis, other_act, jnp.zeros_like(other_act)).reshape((-1,))\n",
    "        obs = jnp.concatenate([jnp.zeros(self.num_types,), jnp.zeros((2,)), other_act, other_vis[None,]], axis=-1)\n",
    "        return obs, state\n",
    "    \n",
    "    def action_space(\n",
    "        self, params: Optional[EnvParams] = None\n",
    "    ) -> spaces.Discrete:\n",
    "        \"\"\"Action space of the environment.\"\"\"\n",
    "        return spaces.Discrete(self.num_types)\n",
    "\n",
    "    def observation_space(self, params: EnvParams) -> spaces.Box:\n",
    "        \"\"\"Observation space of the environment.\"\"\"\n",
    "        return spaces.Box(jnp.zeros((self.num_types*(self.num_agents+1)+3,)), jnp.ones((self.num_types*(self.num_agents+1)+3,)), (self.num_types*(self.num_agents+1)+3,), dtype=jnp.float32)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SHORT SEQUENCE\n",
    "\n",
    "config = {\n",
    "    \"LR\": 2.5e-4,\n",
    "    \"NUM_ENVS\": 16,\n",
    "    \"NUM_STEPS\": 128,\n",
    "    \"TOTAL_TIMESTEPS\": 8e6,\n",
    "    \"UPDATE_EPOCHS\": 2,\n",
    "    \"NUM_MINIBATCHES\": 8,\n",
    "    \"GAMMA\": 0.99,\n",
    "    \"GAE_LAMBDA\": 0.95,\n",
    "    \"CLIP_EPS\": 0.2,\n",
    "    \"ENT_COEF\": 0.03,\n",
    "    \"VF_COEF\": 0.5,\n",
    "    \"MAX_GRAD_NORM\": 0.5,\n",
    "    \"ANNEAL_LR\": True,\n",
    "    \"DEBUG\": False,\n",
    "    \"CONTINUOUS\": False,\n",
    "    \"UPDATE_VIS_PROB\": False,\n",
    "}\n",
    "\n",
    "num_types = 3\n",
    "num_trials = 4\n",
    "max_trial_len = 24\n",
    "\n",
    "env = MetaMemSeq(\n",
    "    num_types=num_types,\n",
    "    num_trials=num_trials,\n",
    "    max_trial_len=max_trial_len,\n",
    ")\n",
    "\n",
    "config[\"ENV\"] = env\n",
    "config[\"ENV_PARAMS\"] = EnvParams()\n",
    "\n",
    "jit_train = jax.jit(jax.vmap(make_train(config)))\n",
    "rng = jax.random.PRNGKey(64)\n",
    "rngs = jax.random.split(rng, num_seeds)\n",
    "\n",
    "outs_short = jit_train(rngs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LONG SEQUENCE\n",
    "\n",
    "config = {\n",
    "    \"LR\": 2.5e-4,\n",
    "    \"NUM_ENVS\": 16,\n",
    "    \"NUM_STEPS\": 128,\n",
    "    # \"TOTAL_TIMESTEPS\": 1e6,\n",
    "    \"TOTAL_TIMESTEPS\": 8e6,\n",
    "    \"UPDATE_EPOCHS\": 2,\n",
    "    \"NUM_MINIBATCHES\": 8,\n",
    "    \"GAMMA\": 0.99,\n",
    "    \"GAE_LAMBDA\": 0.95,\n",
    "    \"CLIP_EPS\": 0.2,\n",
    "    \"ENT_COEF\": 0.03,\n",
    "    \"VF_COEF\": 0.5,\n",
    "    \"MAX_GRAD_NORM\": 0.5,\n",
    "    \"ANNEAL_LR\": True,\n",
    "    \"DEBUG\": False,\n",
    "    \"CONTINUOUS\": False,\n",
    "    \"UPDATE_VIS_PROB\": False,\n",
    "}\n",
    "\n",
    "num_types = 3\n",
    "num_trials = 32\n",
    "max_trial_len = 24\n",
    "\n",
    "env = MetaMemSeq(\n",
    "    num_types=num_types,\n",
    "    num_trials=num_trials,\n",
    "    max_trial_len=max_trial_len,\n",
    ")\n",
    "\n",
    "config[\"ENV\"] = env\n",
    "config[\"ENV_PARAMS\"] = EnvParams()\n",
    "\n",
    "jit_train = jax.jit(jax.vmap(make_train(config)))\n",
    "rng = jax.random.PRNGKey(64)\n",
    "rngs = jax.random.split(rng, num_seeds)\n",
    "\n",
    "outs_long = jit_train(rngs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_eval_params_and_return(rng, params, env_params):\n",
    "    rng, _rng = jax.random.split(rng)\n",
    "    last_obs, env_state = init_env.reset(_rng, env_params)\n",
    "\n",
    "    init_state = ScannedRNN.initialize_carry(1, 256)\n",
    "\n",
    "    # COLLECT TRAJECTORIES\n",
    "    def _env_step(runner_state, unused):\n",
    "        prev_env_state, env_params, last_obs, last_done, prev_hstate, rng, ever_done, running_r, running_hstate, running_obs = runner_state\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "\n",
    "        # SELECT ACTION\n",
    "        ac_in = (last_obs[None, None, :], last_done[None, None])\n",
    "        hstate, pi, value = network.apply(params, prev_hstate, ac_in)\n",
    "        action = pi.sample(seed=_rng).squeeze(0).squeeze(0)\n",
    "\n",
    "        # STEP ENV\n",
    "        rng, _rng = jax.random.split(rng)\n",
    "        obsv, env_state, reward, done, info = init_env.step(\n",
    "            _rng, prev_env_state, action, env_params\n",
    "        )\n",
    "\n",
    "        # ONLY GET REWARDS FROM FIRST META-EPISODE\n",
    "        temp_r = running_r[prev_env_state.trial_num] + (reward >= 0.0) * (1.0 - ever_done)\n",
    "        running_r = running_r.at[prev_env_state.trial_num].set(temp_r)\n",
    "\n",
    "        # ONLY FINAL HIDDEN STATE AND ACTION\n",
    "        replace_hstate = jnp.logical_and((prev_env_state.trial_num == env.num_trials - 2), ~ever_done)\n",
    "        running_hstate = running_hstate * (1.0 - replace_hstate) + hstate * replace_hstate\n",
    "        running_obs = running_obs * (1.0 - replace_hstate) + obsv * replace_hstate\n",
    "\n",
    "        ever_done = jnp.logical_or(done, ever_done)\n",
    "        # transition = (prev_env_state, action, running_r, reward, hstate)\n",
    "        transition = (env_state, action, reward, replace_hstate, running_r, running_hstate, running_obs)\n",
    "        runner_state = (env_state, env_params, obsv, done, hstate, rng, ever_done, running_r, running_hstate, running_obs)\n",
    "        return runner_state, transition\n",
    "\n",
    "    runner_state = (env_state, env_params, last_obs, False, init_state, rng, False, jnp.zeros((env.num_trials,)), init_state, last_obs)\n",
    "    runner_state, traj_batch = jax.lax.scan(_env_step, runner_state, None, env.max_trial_len*env.num_trials+1)\n",
    "    return traj_batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = jax.random.PRNGKey(12)\n",
    "rng, rng_init = jax.random.split(rng)\n",
    "env_params = EnvParams()\n",
    "init_env = MetaMemSeq(\n",
    "    num_types=num_types,\n",
    "    num_trials=32,\n",
    "    max_trial_len=max_trial_len,\n",
    ")\n",
    "rngs = jax.random.split(rng, 128)\n",
    "\n",
    "test_short = jax.jit(jax.vmap(jax.vmap(init_eval_params_and_return, in_axes=(0, None, None)), in_axes=(None, 0, None)))(rngs, outs_short[0][0].params, env_params)\n",
    "test_long = jax.jit(jax.vmap(jax.vmap(init_eval_params_and_return, in_axes=(0, None, None)), in_axes=(None, 0, None)))(rngs, outs_long[0][0].params, env_params)\n",
    "\n",
    "short = test_short[-3][:, :,-1,:].mean(1)\n",
    "long = test_long[-3][:, :,-1,:].mean(1)\n",
    "\n",
    "short_mean = short.mean(0)\n",
    "long_mean = long.mean(0)\n",
    "\n",
    "short_err = short.std(0) / jnp.sqrt(num_seeds)\n",
    "long_err = long.std(0) / jnp.sqrt(num_seeds)\n",
    "\n",
    "plt.plot(short_mean, label=f\"RL$^2$: 4 Trials\", alpha=0.5, linestyle='dashed')\n",
    "plt.fill_between(jnp.arange(32), short_mean - short_err, short_mean + short_err, alpha=0.1)\n",
    "plt.plot(long_mean, label=f\"RL$^2$: 32 Trials\", alpha=0.5, linestyle='dashed')\n",
    "plt.fill_between(jnp.arange(32), long_mean - long_err, long_mean + long_err, alpha=0.1)\n",
    "\n",
    "k = 1\n",
    "scores = j_all_all_trial_scores[:,k]\n",
    "scores = scores.swapaxes(0,1)\n",
    "\n",
    "colors = cm.viridis([j/num_generations for j in range(num_generations)])  # Use the viridis colormap\n",
    "\n",
    "for i, (m, color) in enumerate(zip(scores, colors)):\n",
    "    mean = m.mean(0)\n",
    "    stderr = m.std(0) / jnp.sqrt(5)\n",
    "    if i == 0 or i == 7:\n",
    "        plt.plot(jnp.arange(i*4,i*4+4), mean, label=f\"Generation {i}\", color=color, alpha=1.0)\n",
    "    else:\n",
    "        plt.plot(jnp.arange(i*4,i*4+4), mean, color=color, alpha=1.0)\n",
    "    plt.fill_between(jnp.arange(i * 4, i * 4 + 4), mean - stderr, mean + stderr, alpha=0.1, color=color)\n",
    "\n",
    "\n",
    "plt.hlines(real_ev[k], 0, 32, label=\"Noisy Oracle\", color=\"black\", linestyle=\"--\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"Trial\")\n",
    "plt.ylabel(\"Return\")\n",
    "plt.title(\"In-Context Accumulation Memory Sequence\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "\n",
    "for j, k in enumerate([0, 2, 6]):\n",
    "    scores = merged_all_trial_scores[:,k]\n",
    "    colors = [\"blue\", \"orange\", \"green\"]\n",
    "    scores = scores.swapaxes(0,1)\n",
    "    for i, m in enumerate(scores):\n",
    "        mean = m.mean(0)\n",
    "        stderr = m.std(0) / jnp.sqrt(5)\n",
    "        if i == 0:\n",
    "            plt.plot(jnp.arange(i*4,i*4+4), mean, label=f\"Noisy Oracle Prob: {oracle_probs[k]:.2f}\", alpha=1.0, color=colors[j])\n",
    "        else:\n",
    "            plt.plot(jnp.arange(i*4,i*4+4), mean, alpha=1.0, color=colors[j])\n",
    "        plt.fill_between(jnp.arange(i * 4, i * 4 + 4), mean - stderr, mean + stderr, alpha=0.1, color=colors[j])\n",
    "\n",
    "\n",
    "plt.legend()\n",
    "plt.ylabel(\"Return\")\n",
    "plt.xlabel(\"Trial\")\n",
    "plt.title(\"In-Context Accumulation Memory Sequence\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}