imp_marl/environments/new_imp_env_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:32:23.951893Z",
     "start_time": "2023-06-20T15:32:23.929337Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "from imp_env import ImpEnv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Description of the IMP environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this simple tutorial, we are going to create an IMP environment for an engineering system that deteriorates over time.\n",
    "\n",
    "The damage condition is quantified through four discrete bins: [no_damage, minor_damage, major_damage, failure]. The exact damage is not known and the condition is described by $p$, which is equal to the probability of being in any of the four damage bins . \n",
    "\n",
    "At the initial step, all components damage probabilities are set up as:\n",
    "\n",
    "$p_0=[0.9, 0.08, 0.02, 0]$\n",
    "\n",
    "The system consists of $n$ components, each of them deteriorating according to the following transition model:\n",
    "\n",
    "$$\n",
    "T_0 = \\begin{bmatrix}\n",
    "0.9 & 0.05 & 0.03 & 0.02\\\\\n",
    "0 & 0.9 & 0.06 & 0.04\\\\\n",
    "0 & 0 & 0.94 & 0.06\\\\\n",
    "0 & 0 & 0 & 1\\\\\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "The failure probability of a component is equal to the probability of being in the last damage bin.\n",
    "\n",
    "$p_{f_{comp}} = p(last_{bin})$\n",
    "\n",
    "The failure event is defined as a series systems and the system failure probability can be formulated as:\n",
    "\n",
    "$p_{f_{sys}} = 1 - \\prod_{n_{comp}}[(1-p_{f_{comp}})]$\n",
    "\n",
    "If an inspection is conducted, an indication from two possible outcomes (damage, no_damage) is collected. The inspection model is defined conditional on the damage condition:\n",
    "\n",
    "$p(i_{\\text{no_damage}}|\\text{damage_size})=[0.1, 0.3, 0.6, 0.9]$\n",
    "\n",
    "$p(i_{damage}|\\text{damage_size})=[0.9, 0.7, 0.4, 0.1]$\n",
    "\n",
    "An agents can either:\n",
    "- Do-nothing (action=0): the component damage transitions according to the deterioration model described\n",
    "- Inspect (action=1): the component damage transitions according to the deterioration model described, but the inspection outcome is taken into account in order to update the damage probabilities\n",
    "- Repair (action=2): the component damaga probabilities reset to the initial damage probabilities\n",
    "\n",
    "In the reward model, the action:\n",
    "- Do-nothing (action=0) does not cost anything\n",
    "- Inspect (action=1) costs 1 unit (reward = -1)\n",
    "- Repair (action=2): costs 10 units (reward = -10)\n",
    "\n",
    "Additionally, the system risk is penalised at every time step:\n",
    "- Consequences of a system failure: c_f = 100 (reward = -100)\n",
    "$Risk = p_{{f}_{sys}} \\cdot c_f$\n",
    "\n",
    "The objective is the maxisimisation of the expected sum of discounted rewards."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Steps we will follow to create an IMP environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To implement a new environment, we can go through the following steps:\n",
    "\n",
    "1. Define initial damage probabilities, transition, and inspection models. Alternatively, you can define these models as explained in the tutorial imp_env/pomdp_models/generate_transitions.ipynb. Once they are stored, you can then directly jump to step 2.\n",
    "2. Create a new class from the interface ImpEnv => `class NewIMPenv(ImpEnv):`\n",
    "3. Set up the initialisation of the class including, for instance, the number of components, discount factor, and any other necessary variables => `def __init__(self, config=None):`\n",
    "4. Implement the reset method. This function will reset the episode to the intial stage => `def reset():`\n",
    "5. Implement the step method. This function models the dynamics of the environment at each time step: given the agents' actions, the damage condition probabilities transitions, and a reward is collected => `def step(self, action: dict):`\n",
    "6. Implement additionally required methods. We can define more methods if they are necessary for modelling the dynamics of the environment. In this example, we will need: (i) a transition function, (ii) an immediate_reward function to compute the rewards at each time step, and (iii) a pf_sys function that estimates the system failure probability given the components failure probabilities. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Defining transition and inspection models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initial damage probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:32:48.047506Z",
     "start_time": "2023-06-20T15:32:48.030023Z"
    }
   },
   "outputs": [],
   "source": [
    "initial_damage_proba = np.array([0.9, 0.08, 0.02, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transition model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:32:50.123520Z",
     "start_time": "2023-06-20T15:32:50.108547Z"
    }
   },
   "outputs": [],
   "source": [
    "transition_model = np.zeros((3, 4, 4))\n",
    "\n",
    "# Do-nothing action\n",
    "transition_model[0,0,0] = 0.9\n",
    "transition_model[0,0,1] = 0.05\n",
    "transition_model[0,0,2] = 0.03\n",
    "transition_model[0,0,3] = 0.02\n",
    "\n",
    "transition_model[0,1,1] = 0.9\n",
    "transition_model[0,1,2] = 0.06\n",
    "transition_model[0,1,3] = 0.04\n",
    "\n",
    "transition_model[0,2,2] = 0.94\n",
    "transition_model[0,2,3] = 0.06\n",
    "\n",
    "transition_model[0,3,3] = 1\n",
    "\n",
    "transition_model[1] = transition_model[0]\n",
    "\n",
    "transition_model[2] = np.tile(initial_damage_proba, (4,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:32:55.260445Z",
     "start_time": "2023-06-20T15:32:55.151903Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.9 , 0.08, 0.02, 0.  ])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Transition example ##\n",
    "trans_damage = initial_damage_proba.dot(transition_model[2])\n",
    "trans_damage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inspection model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:32:57.875381Z",
     "start_time": "2023-06-20T15:32:57.803875Z"
    }
   },
   "outputs": [],
   "source": [
    "inspection_model = np.zeros((3, 4, 2))\n",
    "\n",
    "inspection_model[0, :, 0] = np.ones(4)\n",
    "inspection_model[2, :, 0] = np.ones(4)\n",
    "\n",
    "inspection_model[1, :, 1] = np.array([0.1, 0.3, 0.6, 0.9])\n",
    "inspection_model[1, :, 0] = 1 - inspection_model[1, :, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:32:58.807807Z",
     "start_time": "2023-06-20T15:32:58.786390Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.92677346, 0.06407323, 0.00915332, 0.        ])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Bayesian updating example ##\n",
    "bayes_upd = trans_damage*inspection_model[1, :, 0] \n",
    "bayes_norm = bayes_upd / np.sum(bayes_upd)\n",
    "bayes_norm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) Constructing your IMP environment class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:33:03.172420Z",
     "start_time": "2023-06-20T15:33:03.107177Z"
    }
   },
   "outputs": [],
   "source": [
    "class NewIMPenv(ImpEnv):\n",
    "    \n",
    "    ## 3) Set up the initialisation of the class ##\n",
    "    def __init__(self, config=None):\n",
    "        \n",
    "        if config is None:\n",
    "            config = {\"n_comp\": 2,\n",
    "                      \"discount_reward\": 1}\n",
    "            \n",
    "        self.n_comp = config[\"n_comp\"]\n",
    "        self.discount_reward = config[\"discount_reward\"]\n",
    "        self.ep_length = 20  # Horizon length\n",
    "        self.proba_size = 4  # Damage size bins\n",
    "        self.n_obs_inspection = 2  # Total number of possible information received from inspection (crack detected or not)\n",
    "        self.actions_per_agent = 3\n",
    "        self.agent_list = [\"agent_\" + str(i) for i in range(self.n_comp)]\n",
    "        \n",
    "        self.initial_damage_proba = np.zeros((self.n_comp, self.proba_size))\n",
    "        self.initial_damage_proba[:,:] = trans_damage \n",
    "        self.transition_model = transition_model\n",
    "        self.inspection_model = inspection_model\n",
    "        \n",
    "        self.reset()\n",
    "        \n",
    "    ## 4) Implement the reset method ##\n",
    "    def reset(self):\n",
    "        # We need the following line to seed self.np_random\n",
    "        # super().reset(seed=seed)\n",
    "\n",
    "        # Choose the agent's belief\n",
    "        self.time_step = 0\n",
    "        self.damage_proba = self.initial_damage_proba\n",
    "        self.observations = {}\n",
    "        for i in range(self.n_comp):\n",
    "            self.observations[self.agent_list[i]] = np.concatenate(\n",
    "                (self.damage_proba[i], [self.time_step / self.ep_length]))\n",
    "    \n",
    "    ## 5) Implement the step method ##\n",
    "    def step(self, action: dict):\n",
    "        action_list = np.zeros(self.n_comp, dtype=int)\n",
    "        \n",
    "        for i in range(self.n_comp):\n",
    "            action_list[i] = action[self.agent_list[i]]\n",
    "        \n",
    "        inspection, next_proba = self.transition(self.damage_proba, action_list)\n",
    "        \n",
    "        reward_ = self.immediate_reward(next_proba, action_list)\n",
    "        \n",
    "        reward = self.discount_reward ** self.time_step * reward_.item()  # Convert float64 to float\n",
    "\n",
    "        rewards = {}\n",
    "        for i in range(self.n_comp):\n",
    "            rewards[self.agent_list[i]] = reward\n",
    "\n",
    "        self.time_step += 1\n",
    "\n",
    "        self.observations = {}\n",
    "        for i in range(self.n_comp):\n",
    "            self.observations[self.agent_list[i]] = np.concatenate(\n",
    "                (next_proba[i], [self.time_step / self.ep_length]))\n",
    "\n",
    "        self.damage_proba = next_proba\n",
    "\n",
    "        # An episode is done if the agent has reached the target\n",
    "        done = self.time_step >= self.ep_length\n",
    "\n",
    "        # info = {\"belief\": self.beliefs}\n",
    "        return self.observations, rewards, done, inspection\n",
    "    \n",
    "     ## 6) Implement additionally required methods ##\n",
    "    def transition(self, damage_proba, action):\n",
    "        damage_proba_prime = damage_proba\n",
    "        inspection_outcome = np.zeros(self.n_comp, dtype=int)\n",
    "        \n",
    "        for i in range(self.n_comp):\n",
    "            trans_damage = damage_proba[i].dot(self.transition_model[action[i]])\n",
    "            \n",
    "            if action[i] == 1:\n",
    "                insp_nodetec_prob = np.sum(trans_damage*self.inspection_model[1, :, 0])\n",
    "                \n",
    "                if (1-insp_nodetec_prob) < 1e-5:\n",
    "                    inspection_outcome[i] = 0\n",
    "                else:\n",
    "                    ins_dist = np.array([insp_nodetec_prob, 1-insp_nodetec_prob])\n",
    "                    inspection_outcome[i] = np.random.choice(range(0, self.n_obs_inspection), size=None,\n",
    "                                                     replace=True, p=ins_dist)\n",
    "                \n",
    "                trans_damage = trans_damage*self.inspection_model[1, :, inspection_outcome[i]] \n",
    "                trans_damage = trans_damage / np.sum(trans_damage)\n",
    "                \n",
    "            damage_proba_prime[i] = trans_damage\n",
    "        return inspection_outcome, damage_proba_prime\n",
    "    \n",
    "    def immediate_reward(self, damage_proba, action):\n",
    "        reward_system = 0\n",
    "        for i in range(self.n_comp):\n",
    "            if action[i] == 1:\n",
    "                reward_system += -1\n",
    "            elif action[i] == 2:\n",
    "                reward_system += -10\n",
    "        pf_sys = self.pf_sys(damage_proba)\n",
    "        reward_system += pf_sys*(-100)\n",
    "        return reward_system\n",
    "    \n",
    "    def pf_sys(self, damage_proba):\n",
    "        pf = damage_proba[:, -1]\n",
    "        rel_sys = 1\n",
    "        for i in range(self.n_comp):\n",
    "            rel_sys *= (1-pf[i])\n",
    "        return 1-rel_sys\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the newly created IMP environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the number of components and discount factor of your choice:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:33:05.484539Z",
     "start_time": "2023-06-20T15:33:05.465139Z"
    }
   },
   "outputs": [],
   "source": [
    "config = {\n",
    "    'n_comp' : 4,\n",
    "    'discount_reward' : 0.97\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialise an environment `imp_model` from the newly implemented class `NewIMPenv`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:33:06.828675Z",
     "start_time": "2023-06-20T15:33:06.820961Z"
    }
   },
   "outputs": [],
   "source": [
    "imp_model = NewIMPenv(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the initial damage probabilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:33:07.668825Z",
     "start_time": "2023-06-20T15:33:07.648097Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.9 , 0.08, 0.02, 0.  ],\n",
       "       [0.9 , 0.08, 0.02, 0.  ],\n",
       "       [0.9 , 0.08, 0.02, 0.  ],\n",
       "       [0.9 , 0.08, 0.02, 0.  ]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imp_model.initial_damage_proba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assign actions (in this case, do nothing in all components):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:33:08.204027Z",
     "start_time": "2023-06-20T15:33:08.184331Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'agent_0': 0, 'agent_1': 0, 'agent_2': 0, 'agent_3': 0}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "actions_ = {}\n",
    "for k in imp_model.agent_list:\n",
    "    actions_[k] = 0\n",
    "actions_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test the selected action:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:33:08.757392Z",
     "start_time": "2023-06-20T15:33:08.738266Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'agent_0': array([0.81  , 0.117 , 0.0506, 0.0224, 0.05  ]),\n",
       "  'agent_1': array([0.81  , 0.117 , 0.0506, 0.0224, 0.05  ]),\n",
       "  'agent_2': array([0.81  , 0.117 , 0.0506, 0.0224, 0.05  ]),\n",
       "  'agent_3': array([0.81  , 0.117 , 0.0506, 0.0224, 0.05  ])},\n",
       " {'agent_0': -8.66341459329023,\n",
       "  'agent_1': -8.66341459329023,\n",
       "  'agent_2': -8.66341459329023,\n",
       "  'agent_3': -8.66341459329023},\n",
       " False,\n",
       " array([0, 0, 0, 0]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imp_model.step(actions_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulate one episode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-20T15:33:09.290620Z",
     "start_time": "2023-06-20T15:33:09.254600Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step:  1 | actions :  {'agent_0': 1, 'agent_1': 0, 'agent_2': 0, 'agent_3': 0} | reward:  -31.842371043246086\n",
      "Step:  2 | actions :  {'agent_0': 0, 'agent_1': 0, 'agent_2': 1, 'agent_3': 1} | reward:  -31.07608269354173\n",
      "Step:  3 | actions :  {'agent_0': 2, 'agent_1': 2, 'agent_2': 0, 'agent_3': 0} | reward:  -25.210612413432855\n",
      "Step:  4 | actions :  {'agent_0': 0, 'agent_1': 2, 'agent_2': 1, 'agent_3': 0} | reward:  -18.185883121397307\n",
      "Step:  5 | actions :  {'agent_0': 1, 'agent_1': 2, 'agent_2': 1, 'agent_3': 0} | reward:  -19.22923554734219\n",
      "Step:  6 | actions :  {'agent_0': 1, 'agent_1': 2, 'agent_2': 2, 'agent_3': 2} | reward:  -26.93753217047434\n",
      "Step:  7 | actions :  {'agent_0': 2, 'agent_1': 1, 'agent_2': 2, 'agent_3': 0} | reward:  -19.577144588065718\n",
      "Step:  8 | actions :  {'agent_0': 2, 'agent_1': 2, 'agent_2': 1, 'agent_3': 1} | reward:  -18.462641655931147\n",
      "Step:  9 | actions :  {'agent_0': 1, 'agent_1': 0, 'agent_2': 2, 'agent_3': 1} | reward:  -20.679345780675316\n",
      "Step:  10 | actions :  {'agent_0': 2, 'agent_1': 2, 'agent_2': 0, 'agent_3': 1} | reward:  -17.921222004513897\n",
      "Step:  11 | actions :  {'agent_0': 2, 'agent_1': 2, 'agent_2': 2, 'agent_3': 1} | reward:  -23.109150556922998\n",
      "Step:  12 | actions :  {'agent_0': 0, 'agent_1': 2, 'agent_2': 0, 'agent_3': 1} | reward:  -11.267924140133927\n",
      "Step:  13 | actions :  {'agent_0': 1, 'agent_1': 1, 'agent_2': 0, 'agent_3': 0} | reward:  -19.60190833823833\n",
      "Step:  14 | actions :  {'agent_0': 0, 'agent_1': 0, 'agent_2': 0, 'agent_3': 1} | reward:  -21.534025734518476\n",
      "Step:  15 | actions :  {'agent_0': 2, 'agent_1': 0, 'agent_2': 0, 'agent_3': 0} | reward:  -17.634274841445723\n",
      "Step:  16 | actions :  {'agent_0': 2, 'agent_1': 2, 'agent_2': 0, 'agent_3': 1} | reward:  -21.605355124925318\n",
      "Step:  17 | actions :  {'agent_0': 2, 'agent_1': 2, 'agent_2': 2, 'agent_3': 1} | reward:  -19.30468581850743\n",
      "Step:  18 | actions :  {'agent_0': 2, 'agent_1': 1, 'agent_2': 0, 'agent_3': 1} | reward:  -8.857831150587506\n",
      "Step:  19 | actions :  {'agent_0': 2, 'agent_1': 0, 'agent_2': 0, 'agent_3': 1} | reward:  -10.636947964653967\n",
      "Step:  20 | actions :  {'agent_0': 2, 'agent_1': 2, 'agent_2': 1, 'agent_3': 2} | reward:  -17.908275701158686\n",
      "sum_of_rewards:  -400.582450389713\n"
     ]
    }
   ],
   "source": [
    "imp_model.reset()\n",
    "done = False\n",
    "rewards_sum = 0\n",
    "\n",
    "while not done:\n",
    "    actions = {f\"agent_{i}\": random.randint(0,2) for i in range(imp_model.n_comp)}\n",
    "    obs, rewards, done, insp_outcomes = imp_model.step(actions) \n",
    "    rewards_sum += rewards['agent_0']\n",
    "    print('Step: ', imp_model.time_step, '| actions : ', actions, '| reward: ', rewards['agent_0'])\n",
    "    \n",
    "print('sum_of_rewards: ', rewards_sum)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you are ready to create your new environment!!"
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}