nbs/Agents/98_APOBase_commented.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b1b2febe-db6a-462d-aba3-d90ab0d6bb47",
   "metadata": {},
   "source": [
    "# Base (part. Obs.)\n",
    "\n",
    "> Base class containing the core methods of CRLD agents learning under partial observability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "61a99b82-9867-470a-b9ef-a2d9e8beccd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp Agents/POBase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "328aed65-7a4c-44e7-b91c-a23b2948edac",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Imports for the nbdev development environment\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b20e6d62-5f57-4344-ba75-587b9804f740",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8a4aac64-5ced-4616-b210-5a985e5b0f7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import jax\n",
    "import numpy as np\n",
    "import itertools as it\n",
    "\n",
    "import jax.numpy as jnp\n",
    "from jax import grad, jit, vmap\n",
    "from functools import partial\n",
    "\n",
    "from fastcore.utils import *\n",
    "\n",
    "from pyCRLD.Agents.Base import abase\n",
    "from pyCRLD.Utils.Helpers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "283d0d67-3a20-4d4f-9d5e-1547ba792f9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class aPObase(abase):\n",
    "    \"\"\"\n",
    "    Base class for\n",
    "    deterministic policy-average (/ memory mean field) independent (multi-agent) \n",
    "    temporal-difference reinforcement learning with partial observability.\n",
    "\n",
    "    To be used as a base for both, value and policy dynamics.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,\n",
    "                 TransitionTensor,\n",
    "                 RewardTensor,\n",
    "                 ObservationTensor, \n",
    "                 DiscountFactors,\n",
    "                 use_prefactor=False,\n",
    "                 opteinsum=True,\n",
    "                 **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        TransitionTensor : transition model of the environment\n",
    "        RewardTensor : reward model of the environment\n",
    "        DiscountFactors : the agents' discount factors\n",
    "        use_prefactor : use the 1-DiscountFactor prefactor (default: False)\n",
    "        opteinsum : keyword argument to optimize einsum methods (default: True)\n",
    "        \"\"\"\n",
    "        R = jnp.array(RewardTensor)\n",
    "        T = jnp.array(TransitionTensor)\n",
    "        O = jnp.array(ObservationTensor)\n",
    "    \n",
    "        # number of agents\n",
    "        N = R.shape[0]  \n",
    "        assert len(T.shape[1:-1]) == N, \"Inconsistent number of agents\"\n",
    "        assert len(R.shape[2:-1]) == N, \"Inconsistent number of agents\"\n",
    "        assert O.shape[0] == N, \"Inconsistent number of agents\"\n",
    "\n",
    "        # number of actions for each agent        \n",
    "        M = T.shape[1] \n",
    "        assert np.allclose(T.shape[1:-1], M), 'Inconsisten number of actions'\n",
    "        assert np.allclose(R.shape[2:-1], M), 'Inconsisten number of actions'\n",
    "        \n",
    "        # number of states\n",
    "        Z = T.shape[0] \n",
    "        assert T.shape[-1] == Z, 'Inconsisten number of states'\n",
    "        assert R.shape[-1] == Z, 'Inconsisten number of states'\n",
    "        assert R.shape[1] == Z, 'Inconsisten number of states'\n",
    "        assert O.shape[1] == Z, 'Inconsistent number of states'\n",
    "\n",
    "        # number of observations\n",
    "        Q = O.shape[-1]\n",
    "        \n",
    "        self.R, self.T, self.O = R, T, O\n",
    "        self.N, self.M, self.Z, self.Q = N, M, Z, Q\n",
    "        \n",
    "        # discount factors\n",
    "        self.gamma = make_variable_vector(DiscountFactors, N)\n",
    "\n",
    "        # use (1-DiscountFactor) prefactor to have values on scale of rewards\n",
    "        self.pre = 1 - self.gamma if use_prefactor else np.ones(N)        \n",
    "        self.use_prefactor = use_prefactor\n",
    "\n",
    "        # 'load' the other agents actions summation tensor for speed\n",
    "        self.Omega = self._OtherAgentsActionsSummationTensor()\n",
    "        \n",
    "        # state and obs distribution helpers\n",
    "        self.has_last_statdist = False\n",
    "        self._last_statedist = jnp.ones(Z) / Z\n",
    "        self.has_last_obsdist = False\n",
    "        self._last_obsdist = jnp.ones((N, Q)) / Q\n",
    "        \n",
    "        # use optimized einsum method\n",
    "        self.opti = opteinsum  \n",
    "\n",
    "   \n",
    "    # =========================================================================\n",
    "    #   Strategy averaging\n",
    "    # =========================================================================\n",
    "    @partial(jit, static_argnums=0)    \n",
    "    def Xisa(self, X):\n",
    "        \"\"\"\n",
    "        Compute state-action policy given the current observation-action policy\n",
    "        \"\"\"\n",
    "        i = 0; a = 1; s = 2; o = 4  # variables\n",
    "        args = [self.O, [i, s, o], X, [i, o, a], [i, s, a]]\n",
    "        Xisa = jnp.einsum(*args, optimize=self.opti)\n",
    "    \n",
    "        # assert np.allclose(Xisa.sum(-1), 1.0), 'Not a policy. Must not happen!'\n",
    "        return Xisa\n",
    "\n",
    "    @partial(jit, static_argnums=0)\n",
    "    def Tss(self, X):\n",
    "        \"\"\"Compute average transition model Tss given policy X\"\"\"\n",
    "        Xisa = self.Xisa(X)\n",
    "        return abase.Tss(self, Xisa)\n",
    "    \n",
    "    def Bios(self, X):\n",
    "        \"\"\"\n",
    "        Compute 'belief' that environment is in stats s given agent i\n",
    "        observes observation o (Bayes Rule)\n",
    "        \"\"\"\n",
    "        pS = self.statedist(X)\n",
    "        return self._bios(X, pS)\n",
    "    \n",
    "    @partial(jit, static_argnums=0)\n",
    "    def _bios(self, X, pS):\n",
    "        i, s, o = 0, 1, 2 # variables \n",
    "\n",
    "        b = jnp.einsum(self.O, [i,s,o], pS, [s], [i,s,o], optimize=self.opti)\n",
    "        bsum = b.sum(axis=1, keepdims=True)\n",
    "        bsum = bsum + (bsum == 0)  # to avoid dividing by zero\n",
    "        Biso = b / bsum\n",
    "        Bios = jnp.swapaxes(Biso, 1,-1)\n",
    "        \n",
    "        return Bios\n",
    "        \n",
    "    @partial(jit, static_argnums=0)\n",
    "    def fast_Bios(self, X):\n",
    "        \"\"\"\n",
    "        Compute 'belief' that environment is in stats s given agent i\n",
    "        observes observation o (Bayes Rule)\n",
    "        \n",
    "        Unsafe when stationary state distribution is not unique\n",
    "        (i.e., when policies are too extreme)\n",
    "        \"\"\"\n",
    "        i, s, o = 0, 1, 2 # variables \n",
    "        # pS = self.statedist(X) # from full obs base (requires Tss from above)\n",
    "        pS = self._jaxPs(X, self._last_statedist)\n",
    "\n",
    "        b = jnp.einsum(self.O, [i,s,o], pS, [s], [i,s,o], optimize=self.opti)\n",
    "        bsum = b.sum(axis=1, keepdims=True)\n",
    "        bsum = bsum + (bsum == 0)  # to avoid dividing by zero\n",
    "        Biso = b / bsum\n",
    "        Bios = jnp.swapaxes(Biso, 1,-1)\n",
    "        \n",
    "        return Bios\n",
    "    \n",
    "    @partial(jit, static_argnums=0)    \n",
    "    def Tioo(self, X, Bios=None, Xisa=None):\n",
    "        \"\"\"Compute average transition model Tioo, given joint policy X\"\"\"\n",
    "        # For speed up\n",
    "        Bios = self.fast_Bios(X) if Bios is None else Bios\n",
    "        Xisa = self.Xisa(X) if Xisa is None else Xisa\n",
    "        \n",
    "        # variables \n",
    "        # agent i, state s, next state s_, observation o, next obs o', all acts\n",
    "        i = 0; s = 1; s_ = 2; o = 3; o_ = 4; b2d = list(range(5, 5+self.N)) \n",
    "\n",
    "        Y4einsum = list(it.chain(*zip(Xisa,\n",
    "                                      [[s, b2d[a]] for a in range(self.N)])))\n",
    "        \n",
    "        args = [Bios, [i, o, s]] + Y4einsum + [self.T, [s]+b2d+[s_],\n",
    "                self.O, [i, s_, o_], [i, o, o_]]\n",
    "        return jnp.einsum(*args, optimize=self.opti)\n",
    "    \n",
    "    @partial(jit, static_argnums=0)    \n",
    "    def Tioao(self, X, Bios=None, Xisa=None):\n",
    "        \"\"\"Compute average transition model Tioao, given joint policy X\"\"\"\n",
    "        # For speed up\n",
    "        Bios = self.fast_Bios(X) if Bios is None else Bios\n",
    "        Xisa = self.Xisa(X) if Xisa is None else Xisa\n",
    "        \n",
    "        # Variables\n",
    "        # agent i, act a, state s, next state s_, observation o, next obs o_\n",
    "        i = 0; a = 1; s = 2; s_ = 3; o = 4; o_ = 5;\n",
    "        j2k = list(range(6, 6+self.N-1))  # other agents\n",
    "        b2d = list(range(6+self.N-1, 6+self.N-1 + self.N))  # all actions\n",
    "        e2f = list(range(5+2*self.N, 5+2*self.N + self.N-1))  # all other acts\n",
    "\n",
    "        sumsis = [[j2k[l], s, e2f[l]] for l in range(self.N-1)]  # sum inds\n",
    "        otherY = list(it.chain(*zip((self.N-1)*[Xisa], sumsis)))\n",
    "\n",
    "        args = [self.Omega, [i]+j2k+[a]+b2d+e2f,\n",
    "                Bios, [i, o, s]] + otherY + [self.T, [s]+b2d+[s_],\n",
    "                self.O, [i, s_, o_], [i, o, a, o_]]\n",
    "        return jnp.einsum(*args, optimize=self.opti)    \n",
    "    \n",
    "    @partial(jit, static_argnums=0)    \n",
    "    def Rioa(self, X, Bios=None, Xisa=None):\n",
    "        \"\"\"Compute average reward Riosa, given joint policy X \"\"\"\n",
    "        # For speed up\n",
    "        Bios = self.fast_Bios(X) if Bios is None else Bios\n",
    "        Xisa = self.Xisa(X) if Xisa is None else Xisa\n",
    "        \n",
    "        # Variables\n",
    "        # agent i, act a, state s, next state s_, observation o\n",
    "        i = 0; a = 1; s = 2; s_ = 3; o = 4\n",
    "        j2k = list(range(5, 5+self.N-1))  # other agents\n",
    "        b2d = list(range(5+self.N-1, 5+self.N-1 + self.N))  # all actions\n",
    "        e2f = list(range(4+2*self.N, 4+2*self.N + self.N-1))  # all other acts\n",
    " \n",
    "        sumsis = [[j2k[l], s, e2f[l]] for l in range(self.N-1)]  # sum inds\n",
    "        otherY = list(it.chain(*zip((self.N-1)*[Xisa], sumsis)))\n",
    "\n",
    "        args = [self.Omega, [i]+j2k+[a]+b2d+e2f, Bios, [i, o, s]] +\\\n",
    "                otherY + [self.T, [s]+b2d+[s_], self.R, [i, s]+b2d+[s_],\n",
    "                [i, o, a]]\n",
    "\n",
    "        return jnp.einsum(*args, optimize=self.opti)\n",
    "    \n",
    "    @partial(jit, static_argnums=0)        \n",
    "    def Rio(self, X, Bios=None, Xisa=None, Rioa=None):\n",
    "        \"\"\"Compute average reward Rio, given joint policy X\"\"\"       \n",
    "        # For speed up\n",
    "        if Rioa is None:  # Compute Rio from scratch\n",
    "            Bios = self.fast_Bios(X) if Bios is None else Bios\n",
    "            Xisa = self.Xisa(X) if Xisa is None else Xisa\n",
    "            \n",
    "            # Variables\n",
    "            # agent i, state s, next state s_, observation o,  # all actions\n",
    "            i = 0; s = 1; s_ = 2; o = 3; b2d = list(range(4, 4+self.N)) \n",
    "            \n",
    "            Y4einsum = list(it.chain(*zip(Xisa,\n",
    "                                    [[s, b2d[a]] for a in range(self.N)])))\n",
    "            \n",
    "            args = [Bios, [i, o, s]] + Y4einsum + [self.T, [s]+b2d+[s_],\n",
    "                    self.R, [i, s]+b2d+[s_], [i, o]]\n",
    "            return jnp.einsum(*args, optimize=self.opti)\n",
    "        \n",
    "        else:  # Compute Rio based on Rioa (should be faster by factor 20)\n",
    "            i=0; o=1; a=2  # Variables\n",
    "            args = [X, [i, o, a], Rioa, [i, o, a], [i, o]]\n",
    "            return jnp.einsum(*args, optimize=self.opti)\n",
    "\n",
    "    @partial(jit, static_argnums=0)        \n",
    "    def Vio(self, X,\n",
    "            Rio=None, Tioo=None, Bios=None, Xisa=None, Rioa=None,\n",
    "            gamma=None):\n",
    "        \"\"\"Compute average observation values Vio, given joint policy X\"\"\"\n",
    "        gamma = self.gamma if gamma is None else gamma \n",
    "\n",
    "        # For speed up\n",
    "        Bios = self.fast_Bios(X) if Bios is None else Bios\n",
    "        Xisa = self.Xisa(X) if Xisa is None else Xisa\n",
    "        Rio = self.Rio(X, Bios=Bios, Xisa=Xisa, Rioa=Rioa) if Rio is None\\\n",
    "            else Rio\n",
    "        Tioo = self.Tioo(X, Bios=Bios, Xisa=Xisa) if Tioo is None\\\n",
    "            else Tioo\n",
    "        \n",
    "        i = 0; o = 1; op = 2  # Variables\n",
    "        n = np.newaxis\n",
    "        Mioo = np.eye(self.Q)[n,:,:] - gamma[:, n, n] * Tioo\n",
    "        invMioo = jnp.linalg.inv(Mioo)\n",
    "\n",
    "        return self.pre[:,n] * jnp.einsum(invMioo, [i, o, op], Rio, [i, op],\n",
    "                                          [i, o], optimize=self.opti)    \n",
    "\n",
    "    @partial(jit, static_argnums=0)            \n",
    "    def Qioa(self, X, Rioa=None, Vio=None, Tioao=None, Bios=None, Xisa=None,\n",
    "             gamma=None):\n",
    "        gamma = self.gamma if gamma is None else gamma \n",
    "        # For speed up\n",
    "        Rioa = self.Rioa(X, Bios=Bios, Xisa=Xisa) if Rioa is None\\\n",
    "            else Rioa\n",
    "        Vio = self.Vio(X, Bios=Bios, Xisa=Xisa, Rioa=Rioa) if Vio is None\\\n",
    "            else Vio        \n",
    "        Tioao = self.Tioao(X, Bios=Bios, Xisa=Xisa) if Tioao is None\\\n",
    "            else Tioao\n",
    "\n",
    "        nextQioa = jnp.einsum(Tioao, [0,1,2,3], Vio, [0,3], [0,1,2],\n",
    "                             optimize=self.opti)\n",
    "        n = np.newaxis\n",
    "        return self.pre[:,n,n] * Rioa + gamma[:,n,n]*nextQioa    \n",
    "    \n",
    "\n",
    "    # =========================================================================\n",
    "    #   HELPERS\n",
    "    # =========================================================================\n",
    "    @partial(jit, static_argnums=0)            \n",
    "    def Ri(self, X):\n",
    "        \"\"\"Compute average reward Ri, given joint policy X\"\"\" \n",
    "        i, o = 0, 1\n",
    "        return jnp.einsum(self.obsdist(X), [i, o], self.Rio(X), [i, o], [i])\n",
    "    \n",
    "    def obsdist(self, X):\n",
    "        if self.has_last_obsdist:\n",
    "            obsdist =  self._jobsdist(X, self._last_obsdist)\n",
    "        else:\n",
    "            obsdist = jnp.array(self._obsdist(X))\n",
    "            self.has_last_obsdist = True\n",
    "            \n",
    "        self._last_obsdist = obsdist\n",
    "        return obsdist\n",
    "\n",
    "    @partial(jit, static_argnums=0)  \n",
    "    def _jobsdist(self, X, pO0, rndkey=42):\n",
    "        \"\"\"Compute stationary distribution, given joint policy X\"\"\"\n",
    "        Tioo = self.Tioo(X)\n",
    "        Dio = jnp.zeros((self.N, self.Q))\n",
    "        \n",
    "        for i in range(self.N):\n",
    "        \n",
    "            pO = compute_stationarydistribution(Tioo[i])\n",
    "            nrS = jnp.where(pO.mean(0)!=-10, 1, 0).sum()\n",
    "\n",
    "            @jit\n",
    "            def single_dist(pO):\n",
    "                return jnp.max(jnp.where(pO.mean(0)!=-10,\n",
    "                                         jnp.arange(pO.shape[0]), -1))\n",
    "            @jit\n",
    "            def multi_dist(pO):\n",
    "                ix = jnp.argmin(jnp.linalg.norm(pO.T - pO0[i], axis=-1))\n",
    "                return ix\n",
    "            \n",
    "            ix = jax.lax.cond(nrS == 1, single_dist, multi_dist, pO)\n",
    "\n",
    "            Dio = Dio.at[i, :].set(pO[:, ix])\n",
    "\n",
    "        return Dio\n",
    "    \n",
    "    def _obsdist(self, X):\n",
    "        \"\"\"Compute stationary distribution, given joint policy X\"\"\"\n",
    "        Tioo = self.Tioo(X)\n",
    "        Dio = np.zeros((self.N, self.Q))\n",
    "        \n",
    "        for i in range(self.N):\n",
    "            pO = np.array(compute_stationarydistribution(Tioo[i]))\n",
    "        \n",
    "            pO = pO[:, pO.mean(0)!=-10]\n",
    "            if len(pO[0]) == 0:  # this happens when the tollerance can distin.\n",
    "                assert False, 'No _statdist return - must not happen'\n",
    "            elif len(pO[0]) > 1:  # Should not happen, in an ideal world\n",
    "                # sidenote: This means an ideal world is ergodic ;)\n",
    "                print(\"More than 1 state-eigenvector found\")\n",
    "                print(pO.round(2))\n",
    "                nr = len(pO[0])\n",
    "                choice = np.random.randint(nr)\n",
    "                print(\"taking random one: \", choice)\n",
    "                pO = pO[:, choice]\n",
    "                        \n",
    "            Dio[i] = pO.flatten()\n",
    "\n",
    "        return Dio\n",
    "    # ===================\n",
    "    # ======================================================\n",
    "    #   Additional state based averages\n",
    "    # =========================================================================\n",
    "    @partial(jit, static_argnums=0)  \n",
    "    def Tisas(self, X):\n",
    "        \"\"\"Compute average transition model Tisas, given joint policy X\"\"\"      \n",
    "        Xisa = self.Xisa(X)\n",
    "        return super().Tisas(self, Xisa)\n",
    "\n",
    "    @partial(jit, static_argnums=0)  \n",
    "    def Risa(self, X):\n",
    "        \"\"\"Compute average reward Risa, given joint policy X\"\"\"\n",
    "        Xisa = self.Xisa(X)\n",
    "        return super().Risa(Xisa)\n",
    "\n",
    "    @partial(jit, static_argnums=0)  \n",
    "    def Ris(self, X, Risa=None):\n",
    "        \"\"\"Compute average reward Ris, given joint policy X\"\"\" \n",
    "        Xisa = self.Xisa(X)\n",
    "        return super().Ris(Xisa, Risa=Risa)\n",
    "    \n",
    "    @partial(jit, static_argnums=0)  \n",
    "    def Vis(self, X, Ris=None, Tss=None, Risa=None):\n",
    "        \"\"\"Compute average state values Vis, given joint policy X\"\"\"\n",
    "        Xisa = self.Xisa(X)\n",
    "        Ris = self.Ris(X) if Ris is None else Ris\n",
    "        Tss = self.Tss(X) if Tss is None else Tss\n",
    "        return super().Vis(Xisa, Ris=Ris, Tss=Tss, Risa=Risa)\n",
    "\n",
    "    @partial(jit, static_argnums=0)  \n",
    "    def Qisa(self, X, Risa=None, Vis=None, Tisas=None):\n",
    "        \"\"\"Compute average state-action values Qisa, given joint policy X\"\"\"\n",
    "        Xisa = self.Xisa(X)\n",
    "        Risa = self.Risa(X) if Risa is None else Risa\n",
    "        Vis = self.Vis(X) if Vis is None else Vis\n",
    "        Tisas = self.Tisas(X) if Tisas is None else Tisas\n",
    "        return super().Qisa(Xisa, Risa=Risa, Vis=Vis, Tisas=Tisas)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7acec4e2-b3f0-43ac-92e6-b6b3c61cd3ba",
   "metadata": {},
   "source": [
    "## Strategy Averaging\n",
    "Core methods to compute the strategy-average reward-prediction error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c22422dc-086d-4782-8c27-e224de4b9ebc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L99){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Xisa\n",
       "\n",
       ">      aPObase.Xisa (X)\n",
       "\n",
       "Compute state-action policy given the current observation-action policy"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L99){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Xisa\n",
       "\n",
       ">      aPObase.Xisa (X)\n",
       "\n",
       "Compute state-action policy given the current observation-action policy"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Xisa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e1630fc7-e0d0-4203-8925-6f9c1264646a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L111){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tss\n",
       "\n",
       ">      aPObase.Tss (X)\n",
       "\n",
       "Compute average transition model Tss given policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L111){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tss\n",
       "\n",
       ">      aPObase.Tss (X)\n",
       "\n",
       "Compute average transition model Tss given policy X"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Tss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5f002075-4362-4212-85b4-1ee4aa11f2cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L116){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Bios\n",
       "\n",
       ">      aPObase.Bios (X)\n",
       "\n",
       "Compute 'belief' that environment is in stats s given agent i\n",
       "observes observation o (Bayes Rule)"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L116){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Bios\n",
       "\n",
       ">      aPObase.Bios (X)\n",
       "\n",
       "Compute 'belief' that environment is in stats s given agent i\n",
       "observes observation o (Bayes Rule)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Bios)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7bb5dcf0-28d6-4ef3-97ff-cc741254bf1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L158){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tioo\n",
       "\n",
       ">      aPObase.Tioo (X, Bios=None, Xisa=None)\n",
       "\n",
       "Compute average transition model Tioo, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L158){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tioo\n",
       "\n",
       ">      aPObase.Tioo (X, Bios=None, Xisa=None)\n",
       "\n",
       "Compute average transition model Tioo, given joint policy X"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Tioo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d6a3a464-3f28-4093-b62a-070fe162c56c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L176){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tioao\n",
       "\n",
       ">      aPObase.Tioao (X, Bios=None, Xisa=None)\n",
       "\n",
       "Compute average transition model Tioao, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L176){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tioao\n",
       "\n",
       ">      aPObase.Tioao (X, Bios=None, Xisa=None)\n",
       "\n",
       "Compute average transition model Tioao, given joint policy X"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Tioao)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7a796183-c646-43c6-b070-fee5aae94b5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L198){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Rioa\n",
       "\n",
       ">      aPObase.Rioa (X, Bios=None, Xisa=None)\n",
       "\n",
       "Compute average reward Riosa, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L198){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Rioa\n",
       "\n",
       ">      aPObase.Rioa (X, Bios=None, Xisa=None)\n",
       "\n",
       "Compute average reward Riosa, given joint policy X"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Rioa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8766a93a-e508-49b0-9c45-a434376a8a1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L221){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Rio\n",
       "\n",
       ">      aPObase.Rio (X, Bios=None, Xisa=None, Rioa=None)\n",
       "\n",
       "Compute average reward Rio, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L221){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Rio\n",
       "\n",
       ">      aPObase.Rio (X, Bios=None, Xisa=None, Rioa=None)\n",
       "\n",
       "Compute average reward Rio, given joint policy X"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Rio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ef2b0dce-1730-40ac-8745-9d305739e7e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L245){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Vio\n",
       "\n",
       ">      aPObase.Vio (X, Rio=None, Tioo=None, Bios=None, Xisa=None, Rioa=None,\n",
       ">                   gamma=None)\n",
       "\n",
       "Compute average observation values Vio, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L245){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Vio\n",
       "\n",
       ">      aPObase.Vio (X, Rio=None, Tioo=None, Bios=None, Xisa=None, Rioa=None,\n",
       ">                   gamma=None)\n",
       "\n",
       "Compute average observation values Vio, given joint policy X"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Vio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "165a9d2e-2737-4734-a5fe-e012ed024113",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L268){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Qioa\n",
       "\n",
       ">      aPObase.Qioa (X, Rioa=None, Vio=None, Tioao=None, Bios=None, Xisa=None,\n",
       ">                    gamma=None)"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L268){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Qioa\n",
       "\n",
       ">      aPObase.Qioa (X, Rioa=None, Vio=None, Tioao=None, Bios=None, Xisa=None,\n",
       ">                    gamma=None)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Qioa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "61f84874-9399-4273-b00c-436c656e3231",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L289){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Ri\n",
       "\n",
       ">      aPObase.Ri (X)\n",
       "\n",
       "Compute average reward Ri, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L289){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Ri\n",
       "\n",
       ">      aPObase.Ri (X)\n",
       "\n",
       "Compute average reward Ri, given joint policy X"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Ri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "edec4062-704b-48dc-85a1-5eb546b76629",
   "metadata": {},
   "outputs": [],
   "source": [
    "#show_doc(aPObase.obsdist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b71d6540-a873-450b-bd91-51113afab56b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L358){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tisas\n",
       "\n",
       ">      aPObase.Tisas (X)\n",
       "\n",
       "Compute average transition model Tisas, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L358){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Tisas\n",
       "\n",
       ">      aPObase.Tisas (X)\n",
       "\n",
       "Compute average transition model Tisas, given joint policy X"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Tisas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a9c1b815-7905-451d-8eea-b0669813dd1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L364){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Risa\n",
       "\n",
       ">      aPObase.Risa (X)\n",
       "\n",
       "Compute average reward Risa, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L364){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Risa\n",
       "\n",
       ">      aPObase.Risa (X)\n",
       "\n",
       "Compute average reward Risa, given joint policy X"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Risa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1373c1fd-014a-4024-8b3b-7c2db3c262e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L370){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Ris\n",
       "\n",
       ">      aPObase.Ris (X, Risa=None)\n",
       "\n",
       "Compute average reward Ris, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L370){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Ris\n",
       "\n",
       ">      aPObase.Ris (X, Risa=None)\n",
       "\n",
       "Compute average reward Ris, given joint policy X"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Ris)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c3320a94-f9ab-4dba-bd09-752b3a14a07a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L376){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Vis\n",
       "\n",
       ">      aPObase.Vis (X, Ris=None, Tss=None, Risa=None)\n",
       "\n",
       "Compute average state values Vis, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L376){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Vis\n",
       "\n",
       ">      aPObase.Vis (X, Ris=None, Tss=None, Risa=None)\n",
       "\n",
       "Compute average state values Vis, given joint policy X"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Vis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6e97c7c4-9165-4c01-894e-e35395a4c79d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L384){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Qisa\n",
       "\n",
       ">      aPObase.Qisa (X, Risa=None, Vis=None, Tisas=None)\n",
       "\n",
       "Compute average state-action values Qisa, given joint policy X"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/wbarfuss/pyCRLD/blob/main/pyCRLD/Agents/POBase.py#L384){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### aPObase.Qisa\n",
       "\n",
       ">      aPObase.Qisa (X, Risa=None, Vis=None, Tisas=None)\n",
       "\n",
       "Compute average state-action values Qisa, given joint policy X"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(aPObase.Qisa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9054003f-5f60-41f3-bfa0-a37578f10b7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}