{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "title" }, "source": [ "# Hands-on: Contrastive and Counterfactual explanations" ] }, { "cell_type": "markdown", "metadata": { "id": "intro" }, "source": [ "## 1. Introduction\n", "\n", "This notebook uses one small Overcooked trace to compare several styles of explanation for agent behaviour. The emphasis is on the explanatory questions and the shape of the answers: what is the fact, what is the foil, what changes under an intervention, what future opportunity is enabled, and how a policy graph can support telic explanation.\n", "\n", "The original papers often require substantial machinery: learned Generalized Value Functions, visual generative models, learned causal graphs, recurrent models, or large replay-derived policy graphs. Here we keep the Overcooked example deliberately small and transparent. Whenever we simplify a method, the relevant section says exactly which part of the original method is being approximated.\n", "\n", "Workflow:\n", "\n", "1. Generate 1000 out-of-the-box `GreedyHumanModel` episodes on the `cramped_room` layout and store the traces.\n", "2. Generate one small scripted trace on the same layout for the hand-worked explanation example.\n", "3. Select one concrete decision point.\n", "4. Extract a compact explanation state from the raw Overcooked state.\n", "5. Demonstrate simplified, trace-grounded versions of several explanation styles:\n", " - **Contrastive explanation**: compare the chosen action against a foil. (Lin et al. 2021; Zehtabi et al. 2024 for the broader multi-agent optimisation framing.)\n", " - **Foil-state and intervention reasoning**: compare a nearby observed foil state and a forced branch from the same raw state. (Related to Olson et al. 2021, while omitting their visual generative model.)\n", " - **Opportunity-chain illustration**: compare fact vs foil outcome chains using a hand-defined Overcooked causal sketch. (Inspired by Madumal et al. 2020, while omitting their learned AIG/RNN pipeline.)\n", " - **Policy-graph explanations with telic intention analysis (`pgeon`)**: build a small replay-based graph and compute desires/intentions for sequential explanation queries. (Closest to Gimenez-Abalos et al. 2025, but still small-scale.)" ] }, { "cell_type": "markdown", "metadata": { "id": "imports-md" }, "source": [ "## 2. Setup and trace generation\n", "\n", "Imports and environment setup.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:40.556089Z", "iopub.status.busy": "2026-05-25T21:25:40.555917Z", "iopub.status.idle": "2026-05-25T21:25:41.619536Z", "shell.execute_reply": "2026-05-25T21:25:41.619007Z" }, "id": "install-code", "outputId": "c80e341c-9a02-4e35-f674-800ef63ed77d" }, "outputs": [], "source": [ "%pip install --quiet numpy networkx matplotlib scikit-learn pgeon-xai overcooked-ai pygame\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:41.621715Z", "iopub.status.busy": "2026-05-25T21:25:41.621556Z", "iopub.status.idle": "2026-05-25T21:25:42.817696Z", "shell.execute_reply": "2026-05-25T21:25:42.816976Z" }, "id": "imports-code", "outputId": "a07cdbcf-f4d6-4953-f41d-9dc1413ba7df" }, "outputs": [], "source": [ "import copy\n", "import itertools\n", "import os\n", "import pickle\n", "from enum import Enum\n", "from pprint import pprint\n", "\n", "import numpy as np\n", "import networkx as nx\n", "from gymnasium import spaces\n", "from pgeon import PGBasedPolicy, PGBasedPolicyMode, PGBasedPolicyNodeNotFoundMode, PolicyGraph, Predicate\n", "import matplotlib.pyplot as plt\n", "import pandas as pd\n", "\n", "plt.rcParams['figure.figsize'] = (10, 4)\n", "\n", "# overcooked-ai 1.1.0 uses np.Inf\n", "np.Inf = np.inf\n", "\n", "import gymnasium as gym\n", "import overcooked_ai_py.static as overcooked_static\n", "import overcooked_ai_py.data.planners as overcooked_planners_data\n", "import overcooked_ai_py.planning.planners as overcooked_planners_mod\n", "from overcooked_ai_py.mdp.actions import Action\n", "from overcooked_ai_py.agents.agent import AgentPair, GreedyHumanModel\n", "from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv\n", "from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld\n", "\n", "# numpy >= 2.0 compatibility fix: Action.ALL_ACTIONS is a mixed list of tuples\n", "# and a string, so np.random.choice() can no longer build a homogeneous array\n", "# from it. Patch Action.sample to sample an index instead.\n", "from overcooked_ai_py.mdp.actions import Action as _Action\n", "@staticmethod\n", "def _action_sample_fixed(action_probs):\n", " idx = np.random.choice(len(_Action.ALL_ACTIONS), p=action_probs)\n", " return _Action.ALL_ACTIONS[idx]\n", "_Action.sample = _action_sample_fixed" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:42.819466Z", "iopub.status.busy": "2026-05-25T21:25:42.819196Z", "iopub.status.idle": "2026-05-25T21:25:43.359298Z", "shell.execute_reply": "2026-05-25T21:25:43.358671Z" }, "id": "trace-source-code", "outputId": "8a6d84cf-2467-4322-cbaa-046ef92a07ca" }, "outputs": [], "source": [ "OVERCOOKED_LAYOUT = 'cramped_room'\n", "mdp = OvercookedGridworld.from_layout_name(OVERCOOKED_LAYOUT)\n", "raw_env = OvercookedEnv.from_mdp(mdp, horizon=10)\n", "raw_env.reset(regen_mdp=False)\n", "\n", "# Epsilon-greedy wrapper around GreedyHumanModel (for some randomness)\n", "class EpsilonGreedyHumanModel(GreedyHumanModel):\n", " def __init__(self, mlam, epsilon=0.10, **kwargs):\n", " super().__init__(mlam, **kwargs)\n", " self.epsilon = epsilon\n", "\n", " def action(self, state):\n", " if np.random.random() < self.epsilon:\n", " # np.random.choice can't handle ALL_ACTIONS directly (mixed tuples\n", " # and string), so sample an index and look it up instead.\n", " idx = np.random.randint(Action.NUM_ACTIONS)\n", " random_action = Action.ALL_ACTIONS[idx]\n", " return random_action, {'action_probs': self.a_probs_from_action(random_action)}\n", " return super().action(state)\n", "\n", "\n", "GREEDY_NUM_EPISODES = 1000\n", "GREEDY_TRACE_HORIZON = 50\n", "EPSILON = 0.10\n", "RANDOM_SEED = 42 # fix for reproducibility across tutorial attendees\n", "GREEDY_TRACES_PATH = (\n", " f'overcooked_eps{int(EPSILON * 100):02d}_greedy_human_'\n", " f'{GREEDY_NUM_EPISODES}_seed{RANDOM_SEED}_traces.pkl'\n", ")\n", "\n", "if os.path.exists(GREEDY_TRACES_PATH):\n", " with open(GREEDY_TRACES_PATH, 'rb') as f:\n", " greedy_traces = pickle.load(f)\n", " print(f'Loaded {GREEDY_TRACES_PATH}')\n", "else:\n", " np.random.seed(RANDOM_SEED)\n", " greedy_trace_env = OvercookedEnv.from_mdp(mdp, horizon=GREEDY_TRACE_HORIZON, info_level=0)\n", " greedy_agent_pair = AgentPair(\n", " EpsilonGreedyHumanModel(greedy_trace_env.mlam, epsilon=EPSILON),\n", " EpsilonGreedyHumanModel(greedy_trace_env.mlam, epsilon=EPSILON),\n", " )\n", "\n", " greedy_traces = greedy_trace_env.get_rollouts(\n", " greedy_agent_pair,\n", " num_games=GREEDY_NUM_EPISODES,\n", " info=False,\n", " )\n", "\n", " with open(GREEDY_TRACES_PATH, 'wb') as f:\n", " pickle.dump(greedy_traces, f)\n", " print(f'Stored {GREEDY_NUM_EPISODES} epsilon-greedy episodes (ε={EPSILON}, seed={RANDOM_SEED}) in {GREEDY_TRACES_PATH}')\n", "\n", "print('Greedy trace keys:', sorted(greedy_traces.keys()))\n", "print('Episode lengths shape:', greedy_traces['ep_lengths'].shape)\n", "print('Episode states shape:', greedy_traces['ep_states'].shape)\n", "\n", "JOINT_ACTION_SCRIPT = [\n", " ((0, -1), (0, 0)), # player 0 moves up to the onion corridor\n", " ((-1, 0), (0, 0)), # player 0 faces west toward the onion dispenser\n", " ('interact', (0, 0)),\n", " ((1, 0), (0, 0)), # move east toward the pot lane\n", " ((0, -1), (0, 0)), # face north toward the pot\n", " ('interact', (0, 0)) # place onion in the pot\n", "]\n", "\n", "raw_trace = [raw_env.state.deepcopy()]\n", "step_records = []\n", "for joint_action in JOINT_ACTION_SCRIPT:\n", " next_state, reward, done, info = raw_env.step(joint_action)\n", " step_records.append({\n", " 'joint_action': joint_action,\n", " 'reward': reward,\n", " 'done': done,\n", " 'state': next_state.deepcopy(),\n", " })\n", " raw_trace.append(next_state.deepcopy())\n", "\n", "print('Generated trace length:', len(raw_trace))\n", "print('Layout shape:', mdp.shape)\n", "print('Pot locations:', mdp.get_pot_locations())\n", "print('Onion dispensers:', mdp.get_onion_dispenser_locations())\n", "print('Dish dispensers:', mdp.get_dish_dispenser_locations())\n", "print('Serving locations:', mdp.get_serving_locations())\n" ] }, { "cell_type": "markdown", "metadata": { "id": "HfLMpjMvZs5J" }, "source": [ "## 2.1 Visualise the scripted trajectory\n", "\n", "Before extracting explanation features, it helps to see the trajectory itself. The frames below are rendered directly from the Overcooked simulator states in `raw_trace`. The title of each frame shows the action that led into that state.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 876 }, "execution": { "iopub.execute_input": "2026-05-25T21:25:43.361197Z", "iopub.status.busy": "2026-05-25T21:25:43.361059Z", "iopub.status.idle": "2026-05-25T21:25:44.277150Z", "shell.execute_reply": "2026-05-25T21:25:44.276528Z" }, "id": "cF_E3CiyZs5J", "outputId": "5183d681-7595-4f7f-be0d-e2fc396715b2" }, "outputs": [], "source": [ "def joint_action_label(joint_action):\n", " if joint_action is None:\n", " return 'initial state'\n", " a0, a1 = joint_action\n", " def one_action_label(action):\n", " if action == 'interact':\n", " return 'INTERACT'\n", " move_names = {\n", " (0, -1): 'UP',\n", " (0, 1): 'DOWN',\n", " (-1, 0): 'LEFT',\n", " (1, 0): 'RIGHT',\n", " (0, 0): 'WAIT',\n", " }\n", " return move_names.get(action, str(action))\n", " return f'P0 {one_action_label(a0)} | P1 {one_action_label(a1)}'\n", "\n", "\n", "def render_overcooked_state_image(state, tile_size=48):\n", " # Render one Overcooked state to a numpy image array\n", " import os\n", " os.environ.setdefault('SDL_VIDEODRIVER', 'dummy')\n", " import pygame\n", " from overcooked_ai_py.visualization.state_visualizer import StateVisualizer\n", "\n", " visualizer = StateVisualizer(\n", " is_rendering_hud=False,\n", " is_rendering_action_probs=False,\n", " tile_size=tile_size,\n", " )\n", " surface = visualizer.render_state(state, mdp.terrain_mtx)\n", " return np.transpose(pygame.surfarray.array3d(surface), (1, 0, 2))\n", "\n", "try:\n", " frame_images = [render_overcooked_state_image(state) for state in raw_trace]\n", " frame_titles = ['t=0\\ninitial state'] + [\n", " f't={i + 1}\\n{joint_action_label(record[\"joint_action\"])}'\n", " for i, record in enumerate(step_records)\n", " ]\n", "\n", " n_frames = len(frame_images)\n", " n_cols = min(4, n_frames)\n", " n_rows = int(np.ceil(n_frames / n_cols))\n", " fig, axes = plt.subplots(n_rows, n_cols, figsize=(3.0 * n_cols, 3.2 * n_rows))\n", " axes = np.array(axes).reshape(-1)\n", "\n", " for ax, image, title in zip(axes, frame_images, frame_titles):\n", " ax.imshow(image)\n", " ax.set_title(title, fontsize=9)\n", " ax.axis('off')\n", "\n", " for ax in axes[n_frames:]:\n", " ax.axis('off')\n", "\n", " plt.tight_layout()\n", " plt.show()\n", "except Exception as e:\n", " print('Trajectory rendering skipped because of pygame error?')\n", " print(f'Error: {type(e).__name__}: {e}')" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:44.278857Z", "iopub.status.busy": "2026-05-25T21:25:44.278723Z", "iopub.status.idle": "2026-05-25T21:25:44.282872Z", "shell.execute_reply": "2026-05-25T21:25:44.282316Z" }, "id": "trace-preview-code", "outputId": "aa5ff1f5-91f7-4750-9f79-e5c4c4f2db2d" }, "outputs": [], "source": [ "for i, state in enumerate(raw_trace):\n", " print(f'--- raw state {i} ---')\n", " pprint(state.to_dict())\n", " print('pot states:', dict(mdp.get_pot_states(state)))\n" ] }, { "cell_type": "markdown", "metadata": { "id": "feature-state-md" }, "source": [ "## 3. Extract an explanation state from Overcooked\n", "\n", "The raw Overcooked state is richer than what we need for the explanation pipeline. We therefore extract a compact state with carefully chosen fields, computed from the real environment.\n", "\n", "In a real scenario, this extraction is usually handled in one of two ways:\n", "\n", "1. Human-defined semantic features:\n", " - The explanation variables are chosen by experts and correspond to interpretable domain concepts.\n", " - This is common in XAI work because it makes explanations directly readable by people.\n", " - Example references: Lin et al. 2021 uses human-provided features in the ESP (Embedding Self-Prediction) model; Madumal et al. 2020 uses causal opportunity-chain variables derived from domain semantics.\n", "\n", "2. Feature extraction from raw observations:\n", " - A pipeline maps raw state or sensory inputs to the compact explanation features.\n", " - This may be a rule-based extractor, a learned perception model, or a hybrid of both.\n", " - Example reference: Olson et al. 2021 constructs counterfactual states from raw game observations using a learned generative model.\n", "\n", "In our notebook, the compact state is a domain-derived feature vector: a human-defined abstraction extracted systematically from the Overcooked state. In a production system, that extractor would be treated as an explicit module and validated to ensure it preserves the decision-relevant information.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:44.284384Z", "iopub.status.busy": "2026-05-25T21:25:44.284272Z", "iopub.status.idle": "2026-05-25T21:25:44.291366Z", "shell.execute_reply": "2026-05-25T21:25:44.290874Z" }, "id": "feature-state-code", "outputId": "a7f3133e-368e-4950-f0b4-3262e33eabe8" }, "outputs": [], "source": [ "ACTIONS = ['INTERACT', 'MOVE_UP', 'MOVE_DOWN', 'MOVE_LEFT', 'MOVE_RIGHT', 'WAIT']\n", "ACTION_TO_INDEX = {name: i for i, name in enumerate(ACTIONS)}\n", "INDEX_TO_ACTION = {i: name for i, name in enumerate(ACTIONS)}\n", "\n", "PROPERTY_WEIGHTS = {\n", " 'progress_to_start_cooking': 1.6,\n", " 'progress_to_serving': 1.4,\n", " 'progress_to_getting_dish': 1.0,\n", " 'collision_risk': -1.2,\n", " 'coordination_cost': -0.8,\n", " 'idleness': -0.7,\n", "}\n", "\n", "def get_player_holding(player):\n", " if player.held_object is None:\n", " return 'nothing'\n", " return player.held_object.name\n", "\n", "\n", "def get_pot_status(overcooked_state):\n", " pot_states = mdp.get_pot_states(overcooked_state)\n", " if pot_states.get('ready'):\n", " return 'ready'\n", " if pot_states.get('cooking'):\n", " return 'cooking'\n", " if pot_states.get('1_items') or pot_states.get('2_items') or pot_states.get('3_items'):\n", " # In Overcooked, a pot with one or more onions is not necessarily fully\n", " # cooking yet. We call this state 'preparing' to avoid overstating the\n", " # domain semantics in the explanations below.\n", " return 'preparing'\n", " return 'empty'\n", "\n", "\n", "def facing_pot(player):\n", " x, y = player.position\n", " dx, dy = player.orientation\n", " return (x + dx, y + dy) in mdp.get_pot_locations()\n", "\n", "\n", "def extract_feature_state(overcooked_state, player_idx=0):\n", " player = overcooked_state.players[player_idx]\n", " teammate = overcooked_state.players[1 - player_idx]\n", " pot_location = mdp.get_pot_locations()[0]\n", " dish_location = mdp.get_dish_dispenser_locations()[0]\n", "\n", " feature_state = {\n", " 'holding': get_player_holding(player),\n", " 'pot_status': get_pot_status(overcooked_state),\n", " 'pot_adjacent': facing_pot(player),\n", " 'service_ready': get_pot_status(overcooked_state) == 'ready',\n", " 'dish_ready': get_player_holding(player) == 'dish' or get_pot_status(overcooked_state) == 'ready',\n", " 'teammate_blocking': abs(player.position[0] - teammate.position[0]) + abs(player.position[1] - teammate.position[1]) <= 1,\n", " 'distance_to_pot': abs(player.position[0] - pot_location[0]) + abs(player.position[1] - pot_location[1]),\n", " 'distance_to_dish': abs(player.position[0] - dish_location[0]) + abs(player.position[1] - dish_location[1]),\n", " 'position': player.position,\n", " 'orientation': player.orientation,\n", " 'raw_timestep': overcooked_state.timestep,\n", " }\n", " return feature_state\n", "\n", "\n", "feature_trace = [extract_feature_state(state) for state in raw_trace]\n", "for i, state in enumerate(feature_trace):\n", " print(f'--- feature state {i} ---')\n", " pprint(state)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 558 }, "execution": { "iopub.execute_input": "2026-05-25T21:25:44.293057Z", "iopub.status.busy": "2026-05-25T21:25:44.292943Z", "iopub.status.idle": "2026-05-25T21:25:45.681284Z", "shell.execute_reply": "2026-05-25T21:25:45.680765Z" }, "id": "xq1cc03a-1K5", "outputId": "afb10c33-9508-461f-8653-8cbc26f2bd2b" }, "outputs": [], "source": [ "FEATURE_KEYS = [\n", " 'holding', 'pot_status', 'pot_adjacent', 'service_ready', 'dish_ready',\n", " 'teammate_blocking', 'distance_to_pot', 'distance_to_dish'\n", "]\n", "CATEGORICAL_FEATURE_KEYS = [\n", " 'holding', 'pot_status', 'pot_adjacent', 'service_ready', 'dish_ready', 'teammate_blocking'\n", "]\n", "NUMERIC_FEATURE_KEYS = ['distance_to_pot', 'distance_to_dish']\n", "\n", "\n", "def state_key(state):\n", " return tuple((key, state[key]) for key in FEATURE_KEYS)\n", "\n", "\n", "def action_bucket(action):\n", " if action == 'interact' or action == getattr(Action, 'INTERACT', 'interact'):\n", " return 'INTERACT'\n", " move_names = {\n", " (0, -1): 'MOVE_UP',\n", " (0, 1): 'MOVE_DOWN',\n", " (-1, 0): 'MOVE_LEFT',\n", " (1, 0): 'MOVE_RIGHT',\n", " (0, 0): 'WAIT',\n", " }\n", " return move_names[action]\n", "\n", "\n", "def replay_state_distance(reference_state, candidate_state):\n", " categorical_distance = sum(\n", " reference_state[key] != candidate_state[key]\n", " for key in CATEGORICAL_FEATURE_KEYS\n", " )\n", " numeric_distance = sum(\n", " abs(reference_state[key] - candidate_state[key])\n", " for key in NUMERIC_FEATURE_KEYS\n", " )\n", " return categorical_distance + numeric_distance\n", "\n", "\n", "def compact_state_summary(state):\n", " return ', '.join(f'{key}={state[key]}' for key in FEATURE_KEYS)\n", "\n", "\n", "def build_greedy_replay_rows(traces):\n", " rows = []\n", " num_episodes = len(traces['ep_lengths'])\n", " for episode_id in range(num_episodes):\n", " episode_length = int(traces['ep_lengths'][episode_id])\n", " usable_steps = min(episode_length, len(traces['ep_states'][episode_id]) - 1)\n", " for timestep in range(usable_steps):\n", " raw_state = traces['ep_states'][episode_id][timestep]\n", " next_raw_state = traces['ep_states'][episode_id][timestep + 1]\n", " joint_action = traces['ep_actions'][episode_id][timestep]\n", " compact_state = extract_feature_state(raw_state)\n", " next_compact_state = extract_feature_state(next_raw_state)\n", " p0_action = action_bucket(joint_action[0])\n", " rows.append({\n", " 'episode_id': episode_id,\n", " 'timestep': timestep,\n", " 'raw_state': raw_state,\n", " 'joint_action': joint_action,\n", " 'p0_action': p0_action,\n", " 'next_raw_state': next_raw_state,\n", " 'state': compact_state,\n", " 'next_state': next_compact_state,\n", " 'reward': float(traces['ep_rewards'][episode_id][timestep]),\n", " 'state_key': state_key(compact_state),\n", " 'next_state_key': state_key(next_compact_state),\n", " 'state_summary': compact_state_summary(compact_state),\n", " 'next_state_summary': compact_state_summary(next_compact_state),\n", " })\n", " return rows\n", "\n", "\n", "def nearest_replay_rows(reference_state, rows=None, require_action=None, exclude_action=None):\n", " rows = greedy_replay_rows if rows is None else rows\n", " scored_rows = []\n", " for row in rows:\n", " if require_action is not None and row['p0_action'] != require_action:\n", " continue\n", " if exclude_action is not None and row['p0_action'] == exclude_action:\n", " continue\n", " scored = dict(row)\n", " scored['distance_to_reference'] = replay_state_distance(reference_state, row['state'])\n", " scored['changed_features'] = feature_changes(reference_state, row['state']) if 'feature_changes' in globals() else ''\n", " scored_rows.append(scored)\n", " if not scored_rows:\n", " return pd.DataFrame()\n", " return pd.DataFrame(scored_rows).sort_values(\n", " ['distance_to_reference', 'episode_id', 'timestep']\n", " ).reset_index(drop=True)\n", "\n", "\n", "def action_frequency_table(rows_df):\n", " counts = rows_df['p0_action'].value_counts().rename_axis('action').reset_index(name='count')\n", " counts['probability'] = (counts['count'] / counts['count'].sum()).round(3)\n", " return counts\n", "\n", "\n", "greedy_replay_rows = build_greedy_replay_rows(greedy_traces)\n", "greedy_replay_df = pd.DataFrame([\n", " {key: row[key] for key in [\n", " 'episode_id', 'timestep', 'p0_action', 'reward', 'state_key',\n", " 'next_state_key', 'state_summary', 'next_state_summary'\n", " ]}\n", " for row in greedy_replay_rows\n", "])\n", "\n", "dataset_summary = pd.DataFrame([\n", " {'quantity': 'episodes', 'value': len(greedy_traces['ep_lengths'])},\n", " {'quantity': 'episode horizon', 'value': greedy_traces['ep_states'].shape[1]},\n", " {'quantity': 'usable transitions', 'value': len(greedy_replay_rows)},\n", " {'quantity': 'unique abstract states', 'value': greedy_replay_df['state_key'].nunique()},\n", " {'quantity': 'unique abstract state-action pairs', 'value': greedy_replay_df[['state_key', 'p0_action']].drop_duplicates().shape[0]},\n", " {'quantity': 'unique abstract transitions', 'value': greedy_replay_df[['state_key', 'p0_action', 'next_state_key']].drop_duplicates().shape[0]},\n", " {'quantity': 'mean sparse return', 'value': round(float(np.mean(greedy_traces['ep_returns'])), 3)},\n", "])\n", "\n", "display(dataset_summary)\n", "print('Distribution of player 0 actions observed in the traces:')\n", "display(action_frequency_table(greedy_replay_df))\n" ] }, { "cell_type": "markdown", "metadata": { "id": "observed-md" }, "source": [ "### 3.1 Choose the observed decision point\n", "\n", "We focus on the state just before the final `interact` with the pot. At that moment the agent is holding an onion, adjacent to the pot, and facing it.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:45.683135Z", "iopub.status.busy": "2026-05-25T21:25:45.683009Z", "iopub.status.idle": "2026-05-25T21:25:45.686293Z", "shell.execute_reply": "2026-05-25T21:25:45.685742Z" }, "id": "observed-code", "outputId": "10f1789a-e999-4f18-9e17-ed88bb2b0cff" }, "outputs": [], "source": [ "OBSERVED_INDEX = 5\n", "OBSERVED_RAW_STATE = raw_trace[OBSERVED_INDEX].deepcopy()\n", "OBSERVED_STATE = copy.deepcopy(feature_trace[OBSERVED_INDEX])\n", "FACT_ACTION = 'INTERACT'\n", "FOIL_ACTION = 'MOVE_DOWN'\n", "\n", "print('Observed raw state index:', OBSERVED_INDEX)\n", "pprint(OBSERVED_RAW_STATE.to_dict())\n", "\n", "print('\\nObserved extracted state:')\n", "pprint(OBSERVED_STATE)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 474 }, "execution": { "iopub.execute_input": "2026-05-25T21:25:45.687704Z", "iopub.status.busy": "2026-05-25T21:25:45.687606Z", "iopub.status.idle": "2026-05-25T21:25:45.937793Z", "shell.execute_reply": "2026-05-25T21:25:45.937204Z" }, "id": "3g3b1Otw-1K6", "outputId": "bcb2a805-7c66-4742-9a56-19aa3bcf3d74" }, "outputs": [], "source": [ "observed_nearest_replay = nearest_replay_rows(OBSERVED_STATE)\n", "observed_min_distance = int(observed_nearest_replay['distance_to_reference'].iloc[0])\n", "observed_support_rows = observed_nearest_replay[\n", " observed_nearest_replay['distance_to_reference'] == observed_min_distance\n", "]\n", "\n", "print(f'Number of replay rows at that distance: {len(observed_support_rows)}')\n", "print('Empirical player-0 action frequencies at the nearest matching replay states:')\n", "display(action_frequency_table(observed_support_rows))\n", "\n", "with pd.option_context('display.max_colwidth', None):\n", " display(observed_support_rows[[\n", " 'episode_id', 'timestep', 'p0_action', 'distance_to_reference', 'state_summary'\n", " ]].head(5))\n" ] }, { "cell_type": "markdown", "metadata": { "id": "analysis-helpers-md" }, "source": [ "## 4. Explanation model and scoring\n", "\n", "This section defines a compact **linear property-scoring model**. It is not a learned RL value model; it is a transparent scoring surrogate that lets us demonstrate contrastive, foil/intervention, and opportunity-chain explanations on one Overcooked state.\n", "\n", "For example, in Lin et al. 2021, the ESP model decomposes action values into Generalised Value Function (GVF) components, weighted by attribution vectors learned via ESP-DQN/integrated gradients. Lin et al. note that the special case **γ = 0** reduces to directly combining features. This notebook uses that *style* of one-step feature combination: manually designed property-prediction functions score each candidate action, and an MSX-style (Minimal Sufficient Explanation) selection rule is applied on top.\n", "\n", "The important limitation is that, for simplicity, our feature predictions are not learned GVFs and the fixed weights are not ESP attributions. The resulting explanation is a sound explanation of this **simplified linear scoring model**, not an artifact for a trained ESP agent." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "execution": { "iopub.execute_input": "2026-05-25T21:25:45.939522Z", "iopub.status.busy": "2026-05-25T21:25:45.939397Z", "iopub.status.idle": "2026-05-25T21:25:45.947637Z", "shell.execute_reply": "2026-05-25T21:25:45.947151Z" }, "id": "analysis-helpers-code" }, "outputs": [], "source": [ "def predict_future_properties(state, action):\n", " props = {name: 0.0 for name in PROPERTY_WEIGHTS}\n", " if action == 'INTERACT':\n", " if state['holding'] == 'onion' and state['pot_status'] == 'empty' and state['pot_adjacent']:\n", " props['progress_to_start_cooking'] += 1.0\n", " if state['holding'] == 'dish' and state['pot_status'] == 'ready' and state['pot_adjacent']:\n", " props['progress_to_serving'] += 0.9\n", " elif action in {'MOVE_UP', 'MOVE_DOWN', 'MOVE_LEFT', 'MOVE_RIGHT'}:\n", " if state['pot_status'] in ('preparing', 'cooking'):\n", " props['progress_to_serving'] += 0.3\n", " if action == 'MOVE_DOWN':\n", " props['coordination_cost'] += 0.2 if state['teammate_blocking'] else 0.0\n", " if state['service_ready'] or state['dish_ready']:\n", " props['progress_to_getting_dish'] += 0.2\n", " elif action == 'WAIT':\n", " props['idleness'] += 1.0\n", " if state['teammate_blocking'] and action != 'WAIT':\n", " props['collision_risk'] += 0.3\n", " props['coordination_cost'] += 0.2\n", " return {k: round(v, 3) for k, v in props.items()}\n", "\n", "def score_action_from_properties(properties, weights=PROPERTY_WEIGHTS):\n", " return round(sum(properties[name] * weights[name] for name in weights), 3)\n", "\n", "\n", "def action_analysis(state):\n", " analysis = {}\n", " for action in ACTIONS:\n", " props = predict_future_properties(state, action)\n", " analysis[action] = {\n", " 'properties': props,\n", " 'score': score_action_from_properties(props),\n", " }\n", " return analysis\n", "\n", "\n", "def choose_action(state):\n", " analysis = action_analysis(state)\n", " action = max(analysis, key=lambda a: analysis[a]['score'])\n", " return action, analysis\n", "\n", "\n", "def minimal_sufficient_contrast(fact_props, foil_props, weights=PROPERTY_WEIGHTS):\n", " \"\"\"\n", " Compute a Minimal Sufficient Explanation (MSX) for the preference of fact over foil.\n", "\n", " Implements the MSX from Lin et al. 2021:\n", " Let P = {components with positive attributed contribution to the preference}\n", " N = {all remaining components}\n", " S(E) = sum of |contribution| for components in E\n", " An MSX is argmin{|E| : E ⊆ P, S(E) > S(N)}.\n", " We select the solution greedily by sorting P descending by |contribution|\n", " and including components until the running total strictly exceeds S(N).\n", " \"\"\"\n", " deltas = []\n", " for key in weights:\n", " contribution = (fact_props[key] - foil_props[key]) * weights[key]\n", " deltas.append((key, round(contribution, 3)))\n", "\n", " positive = [(k, c) for k, c in deltas if c > 0]\n", " negative = [(k, c) for k, c in deltas if c < 0]\n", "\n", " # S(N): total magnitude of negatively contributing components\n", " s_n = round(sum(abs(c) for _, c in negative), 3)\n", "\n", " # Sort P descending, then greedily include until S(E) > S(N)\n", " positive.sort(key=lambda item: abs(item[1]), reverse=True)\n", "\n", " selected = []\n", " running = 0.0\n", " for key, contribution in positive:\n", " selected.append({'property': key, 'contribution': contribution})\n", " running = round(running + contribution, 3)\n", " if running > s_n:\n", " break\n", "\n", " # Edge case: no positive components exist (foil is actually preferred)\n", " if not selected:\n", " all_nonzero = [(k, c) for k, c in deltas if c != 0]\n", " if all_nonzero:\n", " all_nonzero.sort(key=lambda x: abs(x[1]), reverse=True)\n", " selected = [{'property': all_nonzero[0][0], 'contribution': all_nonzero[0][1]}]\n", "\n", " return selected\n", "\n", "\n", "def verbalize_contrast(fact, foil, minimal_contrast):\n", " more = [item['property'] for item in minimal_contrast if item['contribution'] > 0]\n", " less = [item['property'] for item in minimal_contrast if item['contribution'] < 0]\n", " parts = [f'The agent chose {fact} rather than {foil}']\n", " if more:\n", " parts.append('because it predicts more ' + ', '.join(more))\n", " if less:\n", " parts.append('and less ' + ', '.join(less))\n", " return ' '.join(parts) + '.'\n" ] }, { "cell_type": "markdown", "metadata": { "id": "contrastive-md" }, "source": [ "## 5. Contrastive explanation\n", "\n", "**Primary reference**: Lin et al. 2021, *Contrastive Explanations for Reinforcement Learning via Embedded Self-Predictions* (ICLR 2021).\n", "\n", "We start with the contrastive question:\n", "\n", "> Why did the agent choose `INTERACT` rather than `MOVE_DOWN`?\n", "\n", "Here `INTERACT` is the **fact**: what actually happened. `MOVE_DOWN` is the **foil**: the alternative we want to compare against.\n", "\n", "Lin et al. explain action preferences through predicted future properties. In their paper, these properties are learned as **GVFs** (*Generalised Value Functions*): value-like predictors for interpretable signals such as expected progress, danger, resource use, or time in a region. Their **ESP** (*Embedded Self-Prediction*) architecture builds the agent's action values from those GVF predictions, so the explanation components are part of the model's own value computation.\n", "\n", "The paper then uses **MSX** (*Minimal Sufficient Explanation*) to keep the answer short: instead of listing every difference between fact and foil, select a small subset that is sufficient to explain why the fact is preferred.\n", "\n", "This notebook uses hand-crafted one-step Overcooked properties with fixed linear weights. The selection rule below mirrors the MSX idea, but the inputs are simplified hand-picked features rather than learned ESP/GVF attributions. The resulting explanation should be read as an explanation of this simplified scoring model." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:45.949383Z", "iopub.status.busy": "2026-05-25T21:25:45.949238Z", "iopub.status.idle": "2026-05-25T21:25:45.953059Z", "shell.execute_reply": "2026-05-25T21:25:45.952519Z" }, "id": "contrastive-code", "outputId": "51d5bf81-bd56-4d8e-97fd-7c256d9167c3" }, "outputs": [], "source": [ "observed_action, observed_analysis = choose_action(OBSERVED_STATE)\n", "summary = {action: {'score': item['score'], **item['properties']} for action, item in observed_analysis.items()}\n", "pprint(summary)\n", "\n", "fact_props = observed_analysis[FACT_ACTION]['properties']\n", "foil_props = observed_analysis[FOIL_ACTION]['properties']\n", "contrast = minimal_sufficient_contrast(fact_props, foil_props)\n", "contrast_text = verbalize_contrast(FACT_ACTION, FOIL_ACTION, contrast)\n", "\n", "print('\\nMinimal sufficient contrast:')\n", "pprint(contrast)\n", "\n", "print('\\nContrastive explanation:')\n", "print(contrast_text)\n" ] }, { "cell_type": "markdown", "metadata": { "id": "counterfactual-md" }, "source": [ "## 6. Foil-state and intervention reasoning\n", "\n", "**Related reference**: Olson et al. 2021, *Counterfactual State Explanations for Reinforcement Learning Agents via Generative Deep Learning*. Olson et al. define a counterfactual state `s′` as a **minimal perturbation** of `s` such that the agent's policy selects a **different** action `a′ ≠ a`. Their method searches in the latent space of a trained adversarial/Wasserstein autoencoder on raw pixel states.\n", "\n", "This section keeps Olson et al.'s useful question in view: what would have to be different for another action to be selected? However, it does not reproduce their generative method. We use the symbolic Overcooked features already extracted above.\n", "\n", "We inspect three simple comparisons:\n", "\n", "1. **Nearest observed foil state**: a state already in the hand-worked trace where the preferred action differs; from `feature_trace[6]`, the state after the factual `INTERACT`, where pot preparation has begun and the preferred action changes to a movement action.\n", "2. **Nearest observed replay foil**: a similar state from the 1000 `GreedyHumanModel` episodes where the observed action is the requested foil `MOVE_DOWN`. This is empirical nearest-neighbour evidence, not yet a generated counterfactual.\n", "3. **Forward intervention branch**: branch from the observed raw state and force the foil action for one step. This answers *'what state follows the foil action?'* and helps make the foil concrete. In this case, we branch from the same raw state, force `MOVE_DOWN` for one step, and check whether the policy decision changes." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:45.954488Z", "iopub.status.busy": "2026-05-25T21:25:45.954388Z", "iopub.status.idle": "2026-05-25T21:25:45.959941Z", "shell.execute_reply": "2026-05-25T21:25:45.959403Z" }, "id": "counterfactual-code", "outputId": "8ba9927b-0d21-412a-dd47-a6fe0928e4ef" }, "outputs": [], "source": [ "def build_env_from_state(raw_state, horizon=3):\n", " env = OvercookedEnv.from_mdp(mdp, horizon=raw_state.timestep + horizon + 1)\n", " env.reset(regen_mdp=False)\n", " env.state = raw_state.deepcopy()\n", " return env\n", "\n", "\n", "def run_branch_from_raw_state(raw_state, first_joint_action, followup_joint_actions=None, horizon=3):\n", " env = build_env_from_state(raw_state, horizon=horizon)\n", " branch_states = [env.state.deepcopy()]\n", " branch_records = []\n", " actions = [first_joint_action] + list(followup_joint_actions or [])\n", " for joint_action in actions:\n", " next_state, reward, done, info = env.step(joint_action)\n", " branch_records.append({'joint_action': joint_action, 'reward': reward, 'done': done, 'state': next_state.deepcopy()})\n", " branch_states.append(next_state.deepcopy())\n", " if done:\n", " break\n", " return branch_states, branch_records\n", "\n", "\n", "FACTUAL_JOINT_ACTION = ('interact', (0, 0))\n", "FOIL_JOINT_ACTION = ((0, 1), (0, 0))\n", "\n", "factual_branch_states, factual_branch_records = run_branch_from_raw_state(OBSERVED_RAW_STATE, FACTUAL_JOINT_ACTION)\n", "foil_branch_states, foil_branch_records = run_branch_from_raw_state(OBSERVED_RAW_STATE, FOIL_JOINT_ACTION)\n", "\n", "FACTUAL_NEXT_STATE = extract_feature_state(factual_branch_states[1])\n", "COUNTERFACTUAL_STATE = extract_feature_state(foil_branch_states[1])\n", "\n", "print('Factual next state:')\n", "pprint(FACTUAL_NEXT_STATE)\n", "\n", "print('\\nFoil intervention next state:')\n", "pprint(COUNTERFACTUAL_STATE)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 859 }, "execution": { "iopub.execute_input": "2026-05-25T21:25:45.961308Z", "iopub.status.busy": "2026-05-25T21:25:45.961212Z", "iopub.status.idle": "2026-05-25T21:25:46.330517Z", "shell.execute_reply": "2026-05-25T21:25:46.329829Z" }, "id": "counterfactual-ranking-code", "outputId": "b8fc5b81-c50d-445f-cd1e-9ffb65493bf3" }, "outputs": [], "source": [ "# Three useful comparisons for the foil question.\n", "# 1. A later scripted state where the preferred action changes.\n", "# 2. The nearest observed replay state where GreedyHumanModel actually chose the requested foil action.\n", "# 3. A one-step branch where we force the foil action from the same raw state.\n", "\n", "observed_decision_change_state = copy.deepcopy(feature_trace[6])\n", "foil_intervention_state = copy.deepcopy(COUNTERFACTUAL_STATE)\n", "\n", "COMPARISON_KEYS = [\n", " 'holding', 'pot_status', 'pot_adjacent',\n", " 'distance_to_pot', 'distance_to_dish', 'teammate_blocking'\n", "]\n", "\n", "\n", "def best_actions_with_scores(state):\n", " analysis = action_analysis(state)\n", " best_score = max(item['score'] for item in analysis.values())\n", " best_actions = [action for action, item in analysis.items() if item['score'] == best_score]\n", " return best_actions, best_score, analysis\n", "\n", "\n", "def format_best_actions(actions):\n", " return ', '.join(actions) if len(actions) > 1 else actions[0]\n", "\n", "\n", "def feature_changes(reference_state, candidate_state, keys=COMPARISON_KEYS):\n", " changes = []\n", " for key in keys:\n", " before = reference_state[key]\n", " after = candidate_state[key]\n", " if before != after:\n", " changes.append(f'{key}: {before} -> {after}')\n", " return '; '.join(changes) if changes else 'no change in selected features'\n", "\n", "\n", "def feature_distance(reference_state, candidate_state):\n", " return replay_state_distance(reference_state, candidate_state)\n", "\n", "\n", "observed_best_actions, observed_best_score, _ = best_actions_with_scores(observed_decision_change_state)\n", "foil_best_actions, foil_best_score, foil_analysis = best_actions_with_scores(foil_intervention_state)\n", "\n", "nearest_any_replay = nearest_replay_rows(OBSERVED_STATE)\n", "nearest_any_distance = int(nearest_any_replay['distance_to_reference'].iloc[0])\n", "nearest_any_rows = nearest_any_replay[nearest_any_replay['distance_to_reference'] == nearest_any_distance]\n", "\n", "replay_foil_candidates = nearest_replay_rows(OBSERVED_STATE, require_action=FOIL_ACTION)\n", "nearest_replay_foil_distance = int(replay_foil_candidates['distance_to_reference'].iloc[0])\n", "nearest_replay_foil_rows = replay_foil_candidates[\n", " replay_foil_candidates['distance_to_reference'] == nearest_replay_foil_distance\n", "]\n", "nearest_replay_foil = nearest_replay_foil_rows.iloc[0]\n", "nearest_replay_foil_state = nearest_replay_foil['state']\n", "\n", "comparison_rows = [\n", " {\n", " 'type': 'Scripted later state where decision changes',\n", " 'source': 'trace index 6, after factual INTERACT',\n", " 'action evidence': format_best_actions(observed_best_actions),\n", " 'feature_distance': feature_distance(OBSERVED_STATE, observed_decision_change_state),\n", " 'changed_features': feature_changes(OBSERVED_STATE, observed_decision_change_state),\n", " },\n", " {\n", " 'type': f'Nearest observed replay foil ({FOIL_ACTION})',\n", " 'source': f'episode {nearest_replay_foil.episode_id}, timestep {nearest_replay_foil.timestep}',\n", " 'action evidence': nearest_replay_foil.p0_action,\n", " 'feature_distance': nearest_replay_foil_distance,\n", " 'changed_features': feature_changes(OBSERVED_STATE, nearest_replay_foil_state),\n", " },\n", " {\n", " 'type': 'Foil intervention branch',\n", " 'source': 'force MOVE_DOWN from the same raw state',\n", " 'action evidence': format_best_actions(foil_best_actions),\n", " 'feature_distance': feature_distance(OBSERVED_STATE, foil_intervention_state),\n", " 'changed_features': feature_changes(OBSERVED_STATE, foil_intervention_state),\n", " },\n", "]\n", "\n", "with pd.option_context('display.max_colwidth', None):\n", " display(pd.DataFrame(comparison_rows))" ] }, { "cell_type": "markdown", "metadata": { "id": "distal-md" }, "source": [ "## 7. Opportunity-chain illustration and policy decision tree\n", "\n", "**Reference**: Madumal et al. 2020, *Distal Explanations for Model-free Explainable Reinforcement Learning*.\n", "\n", "Madumal et al. develop causal and distal explanations for model-free RL agents. The central object is an **Action Influence Graph (AIG)**: a causal graph whose nodes are state/reward variables and whose edges are labelled by actions. Informally, it represents statements like \"taking action `INTERACT` allows `holding=onion` and `pot_adjacent=True` to influence `pot_status`.\"\n", "\n", "The distal paper adds the idea of **opportunity chains**. These are explanations of the form: action/event A enables B, and B later causes or enables C. This fits Overcooked naturally: putting an onion in the pot enables soup preparation, which later enables serving.\n", "\n", "Here we use a small hand-crafted Overcooked causal sketch to illustrate the central explanatory idea: an action can be explained by the future opportunity it enables. The full Madumal pipeline uses richer replay data, learned/parameterised causal structure, recurrent models, and policy decision trees.\n", "\n", "The toy implementation following features:\n", "- A hand-defined Overcooked graph linking `holding`, `pot_adjacent`, `pot_status`, `service_ready`, and `soup_served`.\n", "- Fact/foil outcome chains with \"enables\" semantics.\n", "- A tiny decision tree fitted on the short trace purely for visualization." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:46.332302Z", "iopub.status.busy": "2026-05-25T21:25:46.332176Z", "iopub.status.idle": "2026-05-25T21:25:47.931474Z", "shell.execute_reply": "2026-05-25T21:25:47.930660Z" }, "id": "distal-code", "outputId": "8d6e90f6-4712-46a1-a8f8-dbb90f162061" }, "outputs": [], "source": [ "# Hand-crafted Action Influence Graph\n", "# Nodes: state variables + reward node\n", "# Edges: encode domain-level enabling relationships.\n", "# In an actual implementation, this graph can be learned from observation\n", "#\n", "# holding --[INTERACT]--> pot_status (onion in hand -> put in pot)\n", "# pot_adjacent --[INTERACT]--> pot_status (must face the pot)\n", "# pot_status --[ANY]------> service_ready (cooking -> soup ready)\n", "# service_ready--[INTERACT]--> soup_served (reward node: serve the soup)\n", "\n", "OC_AIG = nx.DiGraph()\n", "OC_AIG.add_nodes_from(\n", " ['holding', 'pot_adjacent', 'pot_status', 'service_ready', 'soup_served']\n", ")\n", "for n in OC_AIG.nodes():\n", " OC_AIG.nodes[n]['is_reward'] = (n == 'soup_served')\n", "\n", "OC_AIG.add_edge('holding', 'pot_status', action='INTERACT',\n", " condition='holding=onion AND pot_adjacent=True')\n", "OC_AIG.add_edge('pot_adjacent', 'pot_status', action='INTERACT',\n", " condition='holding=onion')\n", "OC_AIG.add_edge('pot_status', 'service_ready', action='ANY',\n", " condition='pot_status in {preparing,cooking,ready}')\n", "OC_AIG.add_edge('service_ready', 'soup_served', action='INTERACT',\n", " condition='holding=dish AND pot_adjacent=True')\n", "\n", "print('Overcooked causal sketch:')\n", "for u, v, d in OC_AIG.edges(data=True):\n", " print(f' {u} --[{d[\"action\"]}]--> {v} ({d[\"condition\"]})')\n", "\n", "REWARD_NODE = 'soup_served'\n", "aig_ancestors = nx.ancestors(OC_AIG, REWARD_NODE) | {REWARD_NODE}\n", "print(f'\\nAIG causal ancestors of reward node \"{REWARD_NODE}\": {sorted(aig_ancestors)}')\n", "\n", "_HOLDING_CODES_DT = {'onion': 0, 'dish': 1, 'nothing': 2, 'soup': 3}\n", "_POT_CODES_DT = {'empty': 0, 'preparing': 1, 'cooking': 2, 'ready': 3}\n", "\n", "def _encode_for_tree(state):\n", " return np.array([\n", " _HOLDING_CODES_DT.get(state['holding'], 2),\n", " _POT_CODES_DT.get(state['pot_status'], 0),\n", " int(state['pot_adjacent']),\n", " int(state['service_ready']),\n", " int(state['dish_ready']),\n", " int(state['teammate_blocking']),\n", " state['distance_to_pot'],\n", " state['distance_to_dish'],\n", " ], dtype=np.int32)\n", "\n", "\n", "def distal_outcome_from_raw_states(raw_states):\n", " outcomes = {\n", " 'soup_started': False,\n", " 'service_possible': False,\n", " 'dish_obtained': False,\n", " 'soup_served': False,\n", " }\n", " for state in raw_states:\n", " player0 = state.players[0]\n", " ps = mdp.get_pot_states(state)\n", " if any(ps.get(k) for k in ('1_items', '2_items', '3_items', 'cooking', 'ready')):\n", " outcomes['soup_started'] = True\n", " if ps.get('ready'):\n", " outcomes['service_possible'] = True\n", " if player0.held_object is not None and player0.held_object.name == 'dish':\n", " outcomes['dish_obtained'] = True\n", " if player0.held_object is not None and player0.held_object.name == 'soup':\n", " outcomes['soup_served'] = True\n", " return outcomes\n", "\n", "\n", "# Opportunity chain with \"enables\" semantics\n", "# Each chain entry is a bare action name or \"action enables event\"\n", "# when the action entails a causal milestone in the AIG.\n", "def opportunity_chain_from_raw_states(raw_states, records):\n", " MOVE_NAMES = {\n", " 'interact': 'interact',\n", " (0, -1): 'move_up',\n", " (0, 1): 'move_down',\n", " (1, 0): 'move_right',\n", " (-1, 0): 'move_left',\n", " }\n", " chain = []\n", " for i, record in enumerate(records):\n", " action_name = MOVE_NAMES.get(record['joint_action'][0], 'wait')\n", " if i + 1 < len(raw_states):\n", " prev_ps = mdp.get_pot_states(raw_states[i])\n", " next_ps = mdp.get_pot_states(raw_states[i + 1])\n", " pot_started = (\n", " not any(prev_ps.get(k) for k in ('1_items', '2_items', 'cooking', 'ready'))\n", " and any(next_ps.get(k) for k in ('1_items', '2_items', 'cooking', 'ready'))\n", " )\n", " if pot_started:\n", " chain.append(f'{action_name} enables soup_started')\n", " else:\n", " chain.append(action_name)\n", " else:\n", " chain.append(action_name)\n", "\n", " outcomes = distal_outcome_from_raw_states(raw_states)\n", " if outcomes['soup_served']:\n", " chain.append('soup_served (reward)')\n", " elif outcomes['service_possible']:\n", " chain.append('service_possible (causal path to reward)')\n", " elif outcomes['soup_started']:\n", " chain.append('soup_started (progress toward reward)')\n", " return chain\n", "\n", "\n", "factual_outcomes = distal_outcome_from_raw_states(factual_branch_states)\n", "foil_outcomes = distal_outcome_from_raw_states(foil_branch_states)\n", "factual_chain = opportunity_chain_from_raw_states(factual_branch_states, factual_branch_records)\n", "foil_chain = opportunity_chain_from_raw_states(foil_branch_states, foil_branch_records)\n", "\n", "print('\\nFactual chain: ' + ' -> '.join(factual_chain))\n", "print('Foil chain: ' + ' -> '.join(foil_chain))\n", "print('\\nFactual outcomes:'); pprint(factual_outcomes)\n", "print('\\nFoil outcomes:'); pprint(foil_outcomes)\n", "\n", "\n", "# Policy decision tree\n", "# In Madumal et al., a decision tree can act as an interpretable policy surrogate:\n", "# it tells us which state variables appear on the decision path for an action.\n", "\n", "from sklearn.tree import DecisionTreeClassifier, export_text\n", "from sklearn.tree import _tree as _sk_tree\n", "\n", "FEATURE_NAMES_DT = [\n", " 'holding', 'pot_status', 'pot_adjacent', 'service_ready',\n", " 'dish_ready', 'teammate_blocking', 'distance_to_pot', 'distance_to_dish',\n", "]\n", "\n", "# Use real GreedyHumanModel traces instead of the short scripted feature_trace\n", "X_trace = np.array([_encode_for_tree(row['state']) for row in greedy_replay_rows])\n", "y_trace = [row['p0_action'] for row in greedy_replay_rows]\n", "\n", "dt_policy = DecisionTreeClassifier(max_depth=3, random_state=0)\n", "dt_policy.fit(X_trace, y_trace)\n", "\n", "print('\\nPolicy decision tree:')\n", "print(export_text(dt_policy, feature_names=FEATURE_NAMES_DT))\n", "\n", "obs_enc = _encode_for_tree(OBSERVED_STATE).reshape(1, -1)\n", "dt_pred = dt_policy.predict(obs_enc)[0]\n", "\n", "node_indicator = dt_policy.decision_path(obs_enc)\n", "dt_path_features = {\n", " FEATURE_NAMES_DT[dt_policy.tree_.feature[nid]]\n", " for nid in node_indicator.indices\n", " if dt_policy.tree_.feature[nid] != _sk_tree.TREE_UNDEFINED\n", "}\n", "\n", "# WHY explanation set: features that are both policy-relevant and causally\n", "# connected to the reward node in the hand-defined Overcooked causal sketch.\n", "N_bar = dt_path_features & aig_ancestors\n", "DISTAL_ACTION = 'INTERACT'\n", "\n", "print('\\nPolicy relevance vs causal relevance:')\n", "print(f' DT prediction for observed state : {dt_pred}')\n", "print(f' Policy-side features : {sorted(dt_path_features)}')\n", "print(f' Causal ancestors of {REWARD_NODE} : {sorted(aig_ancestors)}')\n", "print(f' Intersection : {sorted(N_bar)}')\n", "print()\n", "print('Why explanation:')\n", "if N_bar:\n", " print(f' The agent chose {DISTAL_ACTION} because {sorted(N_bar)} is relevant on both sides:')\n", " print(' it appears in the policy-side explanation and lies on the causal path toward serving soup.')\n", "else:\n", " print(f' No overlap was found between policy-side features and the causal path to {REWARD_NODE}.')\n" ] }, { "cell_type": "markdown", "metadata": { "id": "pgeon-md" }, "source": [ "## 8. Policy-graph explanations (`pgeon`)\n", "\n", "**Primary reference**: Gimenez-Abalos et al. 2025, *Policy Graphs and Intention: answering 'why' and 'how' from a telic perspective* (AAMAS 2025).\n", "\n", "A **policy graph** is a compact behavioural model built from observations. Nodes are discretised states, edges are observed transitions, and the graph stores probabilities such as:\n", "\n", "- `P(a|s)`: how often the agent takes action `a` in state `s`;\n", "- `P(s'|s,a)`: how often action `a` in state `s` leads to state `s'`.\n", "\n", "The `pgeon` library adds a telic layer to that graph. *Telic* means explaining an action through the future situation it helps bring about (generally a goal or a desire). The method represents such hypotheses as **desires** and computes **intention values** over the policy graph.\n", "\n", "In this section, we build the policy graph from the 1000 `GreedyHumanModel` replay episodes, not from the forced foil branch. We still query the graph with the scripted decision point so the explanation remains connected to the hand-worked example above.\n", "\n", "The graph is empirical but still compact: states are the symbolic Overcooked features extracted in this notebook, and low-level actions use the primitive Overcooked action names. This is much closer to the pgeon-xai setup than the earlier tiny scripted graph, while remaining readable.\n", "\n", "**Related work**: Gyevnar et al. 2024 (CEMA) is similar in the way that it generates causal natural-language explanations in multi-agent systems using a probabilistic forward-simulation model and counterfactual effect-size ranking." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "execution": { "iopub.execute_input": "2026-05-25T21:25:47.933367Z", "iopub.status.busy": "2026-05-25T21:25:47.933089Z", "iopub.status.idle": "2026-05-25T21:25:48.064552Z", "shell.execute_reply": "2026-05-25T21:25:48.063986Z" }, "id": "pgeon-helpers-code" }, "outputs": [], "source": [ "pgeon_replay_rows = greedy_replay_rows\n", "pgeon_query_state = copy.deepcopy(OBSERVED_STATE)\n", "TRANSITIONS = {}\n", "ACTION_COUNTS_BY_STATE = {}\n", "pgeon_state_by_key = {}\n", "\n", "for row in pgeon_replay_rows:\n", " key = row['state_key']\n", " next_key = row['next_state_key']\n", " action_name = row['p0_action']\n", "\n", " pgeon_state_by_key.setdefault(key, row['state'])\n", " pgeon_state_by_key.setdefault(next_key, row['next_state'])\n", " ACTION_COUNTS_BY_STATE.setdefault(key, {})\n", " ACTION_COUNTS_BY_STATE[key][action_name] = ACTION_COUNTS_BY_STATE[key].get(action_name, 0) + 1\n", " TRANSITIONS.setdefault((key, action_name), [])\n", " TRANSITIONS[(key, action_name)].append(row['next_state'])\n", "\n", "pgeon_replay_states = list(pgeon_state_by_key.values())\n", "pgeon_query_key = state_key(pgeon_query_state)\n", "\n", "if pgeon_query_key in pgeon_state_by_key:\n", " pgeon_query_graph_state = copy.deepcopy(pgeon_state_by_key[pgeon_query_key])\n", " pgeon_query_resolution = 'state found in traces'\n", "else:\n", " nearest = nearest_replay_rows(pgeon_query_state).iloc[0]\n", " pgeon_query_graph_state = copy.deepcopy(nearest['state'])\n", " pgeon_query_resolution = f'state not found in traces, using nearest replay state from episode {nearest.episode_id} at timestep {nearest.timestep} with distance {int(nearest.distance_to_reference)}'\n", "\n", "HOLDING_CODES = {'onion': 0, 'dish': 1, 'nothing': 2, 'soup': 3}\n", "POT_STATUS_CODES = {'empty': 0, 'preparing': 1, 'cooking': 2, 'ready': 3}\n", "INV_HOLDING_CODES = {v: k for k, v in HOLDING_CODES.items()}\n", "INV_POT_STATUS_CODES = {v: k for k, v in POT_STATUS_CODES.items()}\n", "\n", "\n", "def encode_feature_state(state):\n", " return np.array([\n", " HOLDING_CODES[state['holding']],\n", " POT_STATUS_CODES[state['pot_status']],\n", " int(state['pot_adjacent']),\n", " int(state['service_ready']),\n", " int(state['dish_ready']),\n", " int(state['teammate_blocking']),\n", " state['distance_to_pot'],\n", " state['distance_to_dish'],\n", " ], dtype=np.int32)\n", "\n", "\n", "def decode_feature_state(obs):\n", " if isinstance(obs, dict):\n", " return obs\n", " arr = np.asarray(obs).astype(int).tolist()\n", " return {\n", " 'holding': INV_HOLDING_CODES[arr[0]],\n", " 'pot_status': INV_POT_STATUS_CODES[arr[1]],\n", " 'pot_adjacent': bool(arr[2]),\n", " 'service_ready': bool(arr[3]),\n", " 'dish_ready': bool(arr[4]),\n", " 'teammate_blocking': bool(arr[5]),\n", " 'distance_to_pot': int(arr[6]),\n", " 'distance_to_dish': int(arr[7]),\n", " }\n", "\n", "\n", "class ReplayAction(Enum):\n", " INTERACT = 0\n", " MOVE_UP = 1\n", " MOVE_DOWN = 2\n", " MOVE_LEFT = 3\n", " MOVE_RIGHT = 4\n", " WAIT = 5\n", "\n", "\n", "class HoldingValue(Enum):\n", " ONION = 0\n", " DISH = 1\n", " NOTHING = 2\n", " SOUP = 3\n", "\n", "\n", "class PotStatusValue(Enum):\n", " EMPTY = 0\n", " PREPARING = 1\n", " COOKING = 2\n", " READY = 3\n", "\n", "\n", "class BoolValue(Enum):\n", " FALSE = 0\n", " TRUE = 1\n", "\n", "\n", "class DistanceToPotValue(Enum):\n", " D0 = 0\n", " D1 = 1\n", " D2 = 2\n", " D3 = 3\n", " D4 = 4\n", " D5 = 5\n", "\n", "\n", "class DistanceToDishValue(Enum):\n", " D0 = 0\n", " D1 = 1\n", " D2 = 2\n", " D3 = 3\n", " D4 = 4\n", " D5 = 5\n", "\n", "\n", "class HoldingPredicate: pass\n", "class PotStatusPredicate: pass\n", "class PotAdjacentPredicate: pass\n", "class ServiceReadyPredicate: pass\n", "class DishReadyPredicate: pass\n", "class TeammateBlockingPredicate: pass\n", "class DistanceToPotPredicate: pass\n", "class DistanceToDishPredicate: pass\n", "\n", "PREDICATE_MAP = {\n", " 'holding': (HoldingPredicate, {'onion': HoldingValue.ONION, 'dish': HoldingValue.DISH, 'nothing': HoldingValue.NOTHING, 'soup': HoldingValue.SOUP}),\n", " 'pot_status': (PotStatusPredicate, {'empty': PotStatusValue.EMPTY, 'preparing': PotStatusValue.PREPARING, 'cooking': PotStatusValue.COOKING, 'ready': PotStatusValue.READY}),\n", " 'pot_adjacent': (PotAdjacentPredicate, {False: BoolValue.FALSE, True: BoolValue.TRUE}),\n", " 'service_ready': (ServiceReadyPredicate, {False: BoolValue.FALSE, True: BoolValue.TRUE}),\n", " 'dish_ready': (DishReadyPredicate, {False: BoolValue.FALSE, True: BoolValue.TRUE}),\n", " 'teammate_blocking': (TeammateBlockingPredicate, {False: BoolValue.FALSE, True: BoolValue.TRUE}),\n", " 'distance_to_pot': (DistanceToPotPredicate, {i: DistanceToPotValue[f'D{i}'] for i in range(6)}),\n", " 'distance_to_dish': (DistanceToDishPredicate, {i: DistanceToDishValue[f'D{i}'] for i in range(6)}),\n", "}\n", "\n", "\n", "class ReplayDiscretizer:\n", " def discretize(self, observation):\n", " state = decode_feature_state(observation)\n", " return tuple(Predicate(predicate_class, [mapping[state[key]]]) for key, (predicate_class, mapping) in PREDICATE_MAP.items())\n", " def state_to_str(self, state):\n", " return ' | '.join(str(predicate) for predicate in state)\n", " def str_to_state(self, state):\n", " return state\n", " def nearest_state(self, state):\n", " yield state\n", " for candidate in self.get_predicate_space():\n", " if candidate != state:\n", " yield candidate\n", " def all_actions(self):\n", " return [ReplayAction.INTERACT, ReplayAction.MOVE_UP, ReplayAction.MOVE_DOWN, ReplayAction.MOVE_LEFT, ReplayAction.MOVE_RIGHT, ReplayAction.WAIT]\n", " def get_predicate_space(self):\n", " return [self.discretize(encode_feature_state(state)) for state in pgeon_replay_states]\n", "\n", "\n", "class OvercookedReplayEnv(gym.Env):\n", " metadata = {}\n", " def __init__(self, initial_states=None, horizon=GREEDY_TRACE_HORIZON):\n", " self.initial_states = [copy.deepcopy(state) for state in (initial_states or pgeon_replay_states)]\n", " self.state = copy.deepcopy(self.initial_states[0])\n", " self.horizon = horizon\n", " self.step_count = 0\n", " self.action_space = spaces.Discrete(len(ACTIONS))\n", " self.observation_space = spaces.Box(low=np.array([0, 0, 0, 0, 0, 0, 0, 0]), high=np.array([3, 3, 1, 1, 1, 1, 5, 5]), shape=(8,), dtype=np.int32)\n", " def reset(self, seed=None, options=None):\n", " super().reset(seed=seed)\n", " index = 0 if seed is None else seed % len(self.initial_states)\n", " self.state = copy.deepcopy(self.initial_states[index])\n", " self.step_count = 0\n", " return encode_feature_state(self.state), {}\n", " def step(self, action):\n", " action_name = INDEX_TO_ACTION[int(action)]\n", " next_options = TRANSITIONS.get((state_key(self.state), action_name), [copy.deepcopy(self.state)])\n", " next_state = copy.deepcopy(next_options[self.step_count % len(next_options)])\n", " self.state = next_state\n", " self.step_count += 1\n", " reward = float(self.state['pot_status'] in ('cooking', 'ready'))\n", " terminated = self.step_count >= self.horizon\n", " return encode_feature_state(self.state), reward, terminated, False, {}\n", "\n", "\n", "def build_policy_graph_from_replay(rows, discretizer):\n", " graph = PolicyGraph(OvercookedReplayEnv(), discretizer)\n", " graph.clear()\n", " graph._trajectories_of_last_fit = []\n", " graph._is_fit = False\n", "\n", " node_frequencies = {}\n", " for row in rows:\n", " source = discretizer.discretize(encode_feature_state(row['state']))\n", " target = discretizer.discretize(encode_feature_state(row['next_state']))\n", " action = ACTION_TO_INDEX[row['p0_action']]\n", "\n", " if not graph.has_node(source):\n", " graph.add_node(source, frequency=0)\n", " if not graph.has_node(target):\n", " graph.add_node(target, frequency=0)\n", " graph.nodes[source]['frequency'] += 1\n", " node_frequencies[target] = node_frequencies.get(target, 0) + 1\n", "\n", " if not graph.has_edge(source, target, key=action):\n", " graph.add_edge(source, target, key=action, frequency=0, action=action)\n", " graph[source][target][action]['frequency'] += 1\n", "\n", " for node, frequency in node_frequencies.items():\n", " graph.nodes[node]['frequency'] += frequency\n", "\n", " graph._normalize()\n", " graph._is_fit = True\n", " return graph\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:48.066379Z", "iopub.status.busy": "2026-05-25T21:25:48.066237Z", "iopub.status.idle": "2026-05-25T21:25:51.037165Z", "shell.execute_reply": "2026-05-25T21:25:51.036575Z" }, "id": "pgeon-run-code", "outputId": "6711a578-16c9-4f9e-9d5e-e935c4fad7e2" }, "outputs": [], "source": [ "replay_discretizer = ReplayDiscretizer()\n", "replay_env = OvercookedReplayEnv(initial_states=pgeon_replay_states, horizon=GREEDY_TRACE_HORIZON)\n", "policy_graph = build_policy_graph_from_replay(pgeon_replay_rows, replay_discretizer)\n", "pg_policy = PGBasedPolicy(policy_graph, mode=PGBasedPolicyMode.GREEDY, node_not_found_mode=PGBasedPolicyNodeNotFoundMode.FIND_SIMILAR_NODES)\n", "\n", "query_predicate = replay_discretizer.discretize(encode_feature_state(pgeon_query_graph_state))\n", "scripted_query_predicate = replay_discretizer.discretize(encode_feature_state(pgeon_query_state))\n", "\n", "print('\\nGraph nodes:', policy_graph.number_of_nodes())\n", "print('Graph edges:', policy_graph.number_of_edges())\n", "print('Query-to-node mode:', pgeon_query_resolution)\n", "\n", "print('\\nQuestion 1: What action a would you choose in the resolved query state?')\n", "print(policy_graph.question1(query_predicate))\n", "\n", "print('\\nQuestion 2: Under which states would you perform action a?')\n", "for state in policy_graph.question2(ReplayAction.INTERACT.value)[:5]:\n", " print('-', replay_discretizer.state_to_str(state))\n", "\n", "print('\\nQuestion 3: Why would you not do action INTERACT under state s?')\n", "pprint(policy_graph.question3(query_predicate, ReplayAction.INTERACT))\n", "\n", "print('\\nQuestion 3: When would you not do action MOVE_DOWN under state s?')\n", "pprint(policy_graph.question3(query_predicate, ReplayAction.MOVE_DOWN))" ] }, { "cell_type": "markdown", "metadata": { "id": "pF9-fO4lZs5P" }, "source": [ "### Telic explanations via desires and intentions\n", "\n", "`pgeon` adds a *desire–intention layer* on top of the policy graph.\n", "\n", "- **Desire** `d = ⟨S_d, a_d⟩`: a perform-goal hypothesis proposed by the explainee.\n", " - `S_d` is the state region where the desire can be fulfilled.\n", " - `a_d` is the action expected in that state region.\n", " - Example for this tutorial: `d = ⟨{holding = nothing, pot_adjacent = true, pot_status = preparing}, MOVE_DOWN⟩` — the onion has been placed in the pot, and the agent is ready to move on from the pot tile.\n", "\n", "- **Intention value** `I_d(s)`: the probability, computed on the policy graph, that the agent will bring about desire `d` from state `s`.\n", " - Algorithm 1 initialises all intentions to zero and seeds every desire state with `P(a_d | s)`.\n", " - Algorithm 2 propagates that increment backward through parent states using `P(S' = s | S = p)`.\n", " - If the parent is itself in `S_d`, propagation excludes the fulfilling action `a_d`, using `P(S' = s, A ≠ a_d | S = p)`. This avoids counting repeated fulfilments of the same desire.\n", "\n", "- **Commitment threshold** `C`: a cutoff for when we are willing to attribute the desire as an intention. If `I_d(s) ≥ C`, the state is treated as committed to desire `d`.\n", "\n", "- **Telic query** `ΔI_d(a, s)`: how much taking action `a` changes the expected intention value from state `s`:\n", " `E[I_d(s') | s, a] - I_d(s)`.\n", "\n", "This section implements Algorithms 1 and 2 from Gimenez-Abalos et al. 2025 over the empirical replay graph built above.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-05-25T21:25:51.039098Z", "iopub.status.busy": "2026-05-25T21:25:51.038946Z", "iopub.status.idle": "2026-05-25T21:25:51.332296Z", "shell.execute_reply": "2026-05-25T21:25:51.331638Z" }, "id": "nM9zAYQ8Zs5P", "outputId": "b831ceb6-7926-4029-f3e6-fa386f214b5e" }, "outputs": [], "source": [ "def has_predicate_value(predicate_tuple, predicate_name, value_name):\n", " return any(f'{predicate_name}({value_name})' in str(p) for p in predicate_tuple)\n", "\n", "\n", "def satisfies_after_onion_placed_desire(predicate_tuple):\n", " # S_d: the onion has just been placed in the pot and the agent is ready to move on.\n", " holding_nothing = has_predicate_value(predicate_tuple, 'HoldingPredicate', 'NOTHING')\n", " pot_adjacent = has_predicate_value(predicate_tuple, 'PotAdjacentPredicate', 'TRUE')\n", " pot_preparing = has_predicate_value(predicate_tuple, 'PotStatusPredicate', 'PREPARING')\n", " return holding_nothing and pot_adjacent and pot_preparing\n", "\n", "DESIRE_ACTION = ReplayAction.MOVE_DOWN # a_d\n", "\n", "\n", "def action_value(action):\n", " return action.value if hasattr(action, 'value') else action\n", "\n", "\n", "def p_action_in_node(pg, node, target_action):\n", " # P(a_d | s) from the policy graph.\n", " target_val = action_value(target_action)\n", " for act, prob in pg.question1(node, verbose=False):\n", " if action_value(act) == target_val:\n", " return float(prob)\n", " return 0.0\n", "\n", "\n", "def _action_is_desire(act, desire_action):\n", " return action_value(act) == action_value(desire_action)\n", "\n", "\n", "def _counts_from_edge_attrs(attrs):\n", " if 'actions' in attrs:\n", " return dict(attrs['actions'])\n", " if 'action' in attrs:\n", " return {attrs['action']: attrs.get('frequency', 1)}\n", " return {}\n", "\n", "\n", "def edge_action_counts(pg, src, dst):\n", " edge_data = pg.get_edge_data(src, dst)\n", " if edge_data is None:\n", " return {}\n", "\n", " if pg.is_multigraph():\n", " combined = {}\n", " for attrs in edge_data.values():\n", " for act, cnt in _counts_from_edge_attrs(attrs).items():\n", " combined[act] = combined.get(act, 0) + cnt\n", " return combined\n", "\n", " return _counts_from_edge_attrs(edge_data)\n", "\n", "\n", "def transition_probability(pg, parent, child):\n", " # P(S' = child | S = parent), marginalising over actions.\n", " total = sum(\n", " sum(edge_action_counts(pg, parent, dst).values())\n", " for dst in pg.successors(parent)\n", " )\n", " if total == 0:\n", " return 0.0\n", " return sum(edge_action_counts(pg, parent, child).values()) / total\n", "\n", "\n", "def transition_probability_without_desire_action(pg, parent, child, desire_action):\n", " # P(S' = child, A != a_d | S = parent).\n", " total = sum(\n", " sum(edge_action_counts(pg, parent, dst).values())\n", " for dst in pg.successors(parent)\n", " )\n", " if total == 0:\n", " return 0.0\n", " counts = edge_action_counts(pg, parent, child)\n", " non_desire_count = sum(\n", " cnt for act, cnt in counts.items()\n", " if not _action_is_desire(act, desire_action)\n", " )\n", " return non_desire_count / total\n", "\n", "\n", "def get_node_visit_probs(pg):\n", " counts = {}\n", " for node in pg.nodes():\n", " data = pg.nodes[node]\n", " for attr in ('count', 'n', 'visits', 'visit_count', 'freq',\n", " 'weight', 'probability', 'prob', 'n_visits'):\n", " if attr in data:\n", " counts[node] = float(data[attr])\n", " break\n", "\n", " if len(counts) != pg.number_of_nodes():\n", " counts = {node: 0.0 for node in pg.nodes()}\n", " for node in pg.nodes():\n", " outgoing = sum(\n", " sum(edge_action_counts(pg, node, dst).values())\n", " for dst in pg.successors(node)\n", " )\n", " incoming = sum(\n", " sum(edge_action_counts(pg, src, node).values())\n", " for src in pg.predecessors(node)\n", " )\n", " counts[node] = max(outgoing, incoming)\n", "\n", " total = sum(counts.values())\n", " if total <= 0:\n", " raise ValueError('Policy graph does not expose visit counts or transition counts for P(s).')\n", "\n", " return {node: counts[node] / total for node in pg.nodes()}\n", "\n", "\n", "def telic_delta_from_replay_transition(state, action_name, intentions):\n", " next_states = TRANSITIONS.get((state_key(state), action_name))\n", " if not next_states:\n", " return None\n", "\n", " src = replay_discretizer.discretize(encode_feature_state(state))\n", " expected_next_intention = np.mean([\n", " intentions[replay_discretizer.discretize(encode_feature_state(next_state))]\n", " for next_state in next_states\n", " ])\n", " return float(expected_next_intention - intentions[src])\n", "\n", "\n", "def pgeon_backward_propagation(pg, desire_fn, desire_action, epsilon=1e-8):\n", " \"\"\"\n", " Algorithms 1 and 2 from Gimenez-Abalos et al. 2025.\n", "\n", " Algorithm 1:\n", " I_d(s) <- 0 for every graph state.\n", " For every s in S_d, seed increment <- P(a_d | s).\n", "\n", " Algorithm 2:\n", " Add the increment to I_d(s), then propagate it backward to every parent p.\n", " If p is not in S_d, use P(S'=s | S=p).\n", " If p is in S_d, use P(S'=s, A != a_d | S=p), so propagation does not\n", " cross transitions that would fulfil the desire.\n", " \"\"\"\n", " intentions = {node: 0.0 for node in pg.nodes()}\n", " stack = []\n", "\n", " for node in pg.nodes():\n", " if desire_fn(node):\n", " increment = p_action_in_node(pg, node, desire_action)\n", " stack.append((node, increment))\n", "\n", " while stack:\n", " node, increment = stack.pop()\n", " intentions[node] += increment\n", "\n", " for parent in pg.predecessors(node):\n", " if desire_fn(parent):\n", " parent_to_node_prob = transition_probability_without_desire_action(\n", " pg, parent, node, desire_action\n", " )\n", " else:\n", " parent_to_node_prob = transition_probability(pg, parent, node)\n", "\n", " propagable_increment = parent_to_node_prob * increment\n", " if propagable_increment >= epsilon:\n", " stack.append((parent, propagable_increment))\n", "\n", " return intentions\n", "\n", "\n", "intentions_pgeon = pgeon_backward_propagation(\n", " policy_graph,\n", " satisfies_after_onion_placed_desire,\n", " DESIRE_ACTION,\n", " epsilon=1e-3,\n", ")\n", "\n", "print(f'Desire: d = ')\n", "print()\n", "\n", "query_pred = replay_discretizer.discretize(encode_feature_state(pgeon_query_graph_state))\n", "I_query = intentions_pgeon.get(query_pred, 0.0)\n", "\n", "nonzero_intentions = [(node, value) for node, value in intentions_pgeon.items() if value > 0]\n", "print(f'Non-zero intention states: {len(nonzero_intentions)} / {policy_graph.number_of_nodes()}')\n", "print(f'Max I_d(s): {max(intentions_pgeon.values()):.3f}')\n", "print(f'I_d(query state): {I_query:.6g}')\n", "print('Query resolution:', pgeon_query_resolution)\n", "\n", "print('\\nHighest intention states:')\n", "for node, I_d in sorted(nonzero_intentions, key=lambda item: item[1], reverse=True)[:8]:\n", " sd_tag = ' <- s in S_d' if satisfies_after_onion_placed_desire(node) else ''\n", " short = replay_discretizer.state_to_str(node)[:110]\n", " print(f' I_d = {I_d:.3f} | {short}{sd_tag}')\n", "\n", "C = 0.5\n", "committed_nodes = [n for n, I in intentions_pgeon.items() if I >= C]\n", "print(f'\\nCommitment threshold C = {C}')\n", "print(f' Is the agent committed to I_d in state s according to C?: {I_query >= C}')\n", "\n", "visit_probs = get_node_visit_probs(policy_graph)\n", "P_committed = sum(visit_probs[n] for n in committed_nodes)\n", "\n", "delta_I_interact = telic_delta_from_replay_transition(\n", " pgeon_query_graph_state, 'INTERACT', intentions_pgeon\n", ")\n", "delta_I_move_down = telic_delta_from_replay_transition(\n", " pgeon_query_graph_state, 'MOVE_DOWN', intentions_pgeon\n", ")\n", "\n", "def format_delta(value):\n", " return 'not estimated: action not observed from this replay state' if value is None else f'{value:+.6g}'\n", "\n", "\n", "print('\\nTelic query: delta_I_d(a, s) = E[I_d(s_next | a)] - I_d(s)')\n", "print(' estimated from the empirical replay transitions used to fit the policy graph')\n", "print(f' current I_d(s) = {I_query:.6g}')\n", "print(f' delta_I_d(INTERACT, s) = {format_delta(delta_I_interact)}')\n", "print(f' delta_I_d(MOVE_DOWN, s) = {format_delta(delta_I_move_down)}')" ] }, { "cell_type": "markdown", "metadata": { "id": "MK3eGMfgx-4k" }, "source": [ "Paper references:\n", "- **Lin et al. 2021**: Contrastive Explanations for Reinforcement Learning via Embedded Self-Predictions.\n", "- **Zehtabi et al. 2024**: Contrastive Explanations of Centralized Multi-Agent Optimization Solutions (ICAPS 2024).\n", "- **Olson et al. 2021**: Counterfactual State Explanations for Reinforcement Learning Agents via Generative Deep Learning.\n", "- **Madumal et al. 2020**: Distal Explanations for Model-free Explainable Reinforcement Learning.\n", "- **Gimenez-Abalos et al. 2025**: Policy Graphs and Intention: answering 'why' and 'how' from a telic perspective. (AAMAS 2025)\n", "- **Gyevnar et al. 2023/2024**: Causal Explanations for Sequential Decision-Making in Multi-Agent Systems (CEMA, AAMAS 2024)." ] } ], "metadata": { "colab": { "provenance": [] }, "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.13.12" } }, "nbformat": 4, "nbformat_minor": 0 }