analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8c030f81-5bb9-4e2a-a714-1a98e71b907c",
   "metadata": {},
   "source": [
    "# Animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f76ee5a9d301c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "from ray.rllib.algorithms.algorithm import Algorithm\n",
    "from tqdm import tqdm\n",
    "import wandb\n",
    "import numpy as np\n",
    "\n",
    "def process_observations(observations, current_observation, agent_ids, termination=None):\n",
    "    loc_x = [current_observation[key][0] if key in current_observation else 0 for key in agent_ids]\n",
    "    loc_y = [current_observation[key][1] if key in current_observation else 0 for key in agent_ids]\n",
    "    heading = [current_observation[key][2] if key in current_observation else 0 for key in agent_ids]\n",
    "    if termination:\n",
    "        still_in_the_game = [1 if not termination[key] else 0 for key in agent_ids]\n",
    "    else:\n",
    "        still_in_the_game = [1 for _ in agent_ids]\n",
    "    observations[\"loc_x\"].append(np.array(loc_x))\n",
    "    observations[\"loc_y\"].append(np.array(loc_y))\n",
    "    observations[\"heading\"].append(np.array(heading))\n",
    "    observations[\"still_in_the_game\"].append(np.array(still_in_the_game))\n",
    "\n",
    "    return observations\n",
    "\n",
    "def collect_observations(path_to_checkpoint, env, step_per_time_increment, number_of_steps):\n",
    "    # if checkpoint path doesn't start with /\n",
    "    if path_to_checkpoint[0] != \"/\":\n",
    "        # Login to Weights and Biases\n",
    "        wandb.login()\n",
    "        wandb.init()\n",
    "        # download model from wandb\n",
    "        artifact = wandb.use_artifact(path_to_checkpoint, type='model')\n",
    "        path_to_checkpoint = artifact.download()\n",
    "        wandb.finish()\n",
    "\n",
    "    ray.init(ignore_reinit_error=True)\n",
    "    algo = Algorithm.from_checkpoint(path_to_checkpoint)\n",
    "    \n",
    "    observations = {\"loc_x\": [], \"loc_y\": [], \"heading\": [], \"still_in_the_game\": []}\n",
    "    \n",
    "    # Run simulation\n",
    "    observation, _ = env.reset()\n",
    "    for step_count in tqdm(range(number_of_steps)):\n",
    "        actions = {\n",
    "            key: algo.compute_single_action(\n",
    "                obs_batch, policy_id=\"prey\" if env.agents[key].agent_type == 0 else \"predator\"\n",
    "            ) for key, obs_batch in observation.items()\n",
    "        }\n",
    "        for i in range(step_per_time_increment):\n",
    "            actions = actions if i == 0 else None  # the agent use action once and then the physics do the rest\n",
    "            observation, _, termination, _, _ = env.step(actions)\n",
    "            observations = process_observations(observations, observation, env._agent_ids, termination)\n",
    "\n",
    "    \n",
    "    observations[\"loc_x\"] = np.array(observations[\"loc_x\"])\n",
    "    observations[\"loc_y\"] = np.array(observations[\"loc_y\"])\n",
    "    observations[\"heading\"] = np.array(observations[\"heading\"])\n",
    "    observations[\"still_in_the_game\"] = np.array(observations[\"still_in_the_game\"])\n",
    "    \n",
    "    env.close()\n",
    "    ray.shutdown()\n",
    "    \n",
    "    return observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "856fe490-82ef-4d34-b1dc-f87ef8c97860",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from config import run_config\n",
    "from particle_2d_env import Particle2dEnvironment\n",
    "\n",
    "# Initialize environment\n",
    "env = Particle2dEnvironment(run_config[\"env\"])\n",
    "checkpoint_folder = os.getcwd() + \"/ray_results\" \n",
    "checkpoint_path = checkpoint_folder + \"/PPO_2024-05-24_15-25-16/PPO_Particle2dEnvironment_0f58e_00000_0_2024-05-24_15-25-16/checkpoint_000070\"\n",
    "# If we want to use wandb artefact\n",
    "#artefact_path = \"tcazalet_airo/marl-rllib\"\n",
    "#checkpoint_path =artefact_path +  '/checkpoint_PPO_Particle2dEnvironment_a7f4d_00000:v0'\n",
    "observations = collect_observations(\n",
    "    checkpoint_path,\n",
    "    env, \n",
    "    run_config[\"env\"][\"step_per_time_increment\"], \n",
    "    number_of_steps=400\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc32e38d-7f7b-4612-bc51-78416842c65c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "from visualization import generate_animation_3d\n",
    "\n",
    "ani = generate_animation_3d(observations, env, fps=100, fig_width=8, fig_height=8)\n",
    "\n",
    "HTML(ani.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b0fd8d-2414-4e3b-968e-72be6c21ce55",
   "metadata": {},
   "source": [
    "# Metrics\n",
    "\n",
    "now calculated during training"
   ]
  },
  {
   "cell_type": "raw",
   "id": "37f775fb-948b-4db0-8242-c68ebe51ab2e",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "import os\n",
    "from metrics import calculate_dos\n",
    "\n",
    "# List all checkpoint directories\n",
    "checkpoint_folder = os.getcwd() + \"/ray_results\" + \"/PPO_2024-02-19_21-16-08/PPO_CustomEnvironment_1961a_00000_0_2024-02-19_21-16-08/\"\n",
    "checkpoint_dirs = [d for d in os.listdir(checkpoint_folder) if d.startswith(\"checkpoint_\")]\n",
    "\n",
    "dos_list = []\n",
    "for checkpoint_dir in checkpoint_dirs:\n",
    "    # Construct the path to the current checkpoint\n",
    "    path_to_checkpoint = os.path.join(checkpoint_folder, checkpoint_dir)   # or path_to_checkpoint = checkpoint_folder + \"checkpoint_\" + str(i).zfill(6)\n",
    "    observations = collect_observations(path_to_checkpoint, env)\n",
    "    T, N, D = loc_x.shape[0], loc_x.shape[1], env.grid_diagonal\n",
    "    average_dos = calculate_dos_sum(observations)/(T * N * D)\n",
    "    dos_list.append(average_dos)\n"
   ]
  },
  {
   "cell_type": "raw",
   "id": "09cc1514-6e9f-4649-a593-8df180fc29f5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# plot dos curve\n",
    "plt.plot(dos)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "726b9d2d-06a6-4a4d-a6c2-e8768735e6ea",
   "metadata": {},
   "source": [
    "# Action distribution visualisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d6bec75-eb00-4052-a54b-897064368a4a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from ray.rllib.algorithms.algorithm import Algorithm\n",
    "\n",
    "path_to_checkpoint = os.getcwd() + \"/ray_results\" + \"/PPO_2024-05-22_15-31-55/PPO_Particle2dEnvironment_a7f4d_00000_0_2024-05-22_15-31-55/checkpoint_000005\"\n",
    "\n",
    "algo = Algorithm.from_checkpoint(path_to_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7fa3d0a-0168-4329-85fd-338767e9008e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from ipywidgets import interactive, interact, FloatSlider, VBox, HBox\n",
    "from IPython.display import display\n",
    "import torch\n",
    "\n",
    "# The policy is used to define the functions\n",
    "policy = algo.get_policy(\"prey\")\n",
    "\n",
    "# Calculate vision endpoints\n",
    "left_x = env.max_seeing_distance * np.cos(-env.max_seeing_angle)\n",
    "left_y = env.max_seeing_distance * np.sin(-env.max_seeing_angle)\n",
    "right_x = env.max_seeing_distance * np.cos(env.max_seeing_angle)\n",
    "right_y = env.max_seeing_distance * np.sin(env.max_seeing_angle)\n",
    "\n",
    "def update_plot_observation(r, theta, heading):\n",
    "    fig, ax_obs = plt.subplots()\n",
    "\n",
    "    x = r * np.cos(theta)\n",
    "    y = r * np.sin(theta)\n",
    "    \n",
    "    ax_obs.plot(x, y, 'ro') # Red point at the polar coordinates for the observed fish\n",
    "    ax_obs.plot(0, 0, 'bo') # Blue dot in the middle for the fish observing\n",
    "    ax_obs.set_xlim(-env.stage_size, env.stage_size)\n",
    "    ax_obs.set_ylim(-env.stage_size, env.stage_size)\n",
    "    ax_obs.axhline(0, color='black', linewidth=0.5)\n",
    "    ax_obs.axvline(0, color='black', linewidth=0.5)\n",
    "    ax_obs.set_title(\"setting\")\n",
    "\n",
    "    # Length of the heading line\n",
    "    heading_length = 1  # Adjust as needed\n",
    "    # Calculate the end point of the heading line\n",
    "    heading_x = x + heading_length * np.cos(heading)\n",
    "    heading_y = y + heading_length * np.sin(heading)\n",
    "\n",
    "    # Draw the heading line\n",
    "    ax_obs.plot([x, heading_x], [y, heading_y], 'g-')  # Green line for heading\n",
    "\n",
    "    # Draw the field of vision lines\n",
    "    ax_obs.plot([0, left_x], [0, left_y], 'b--')  # Left vision line\n",
    "    ax_obs.plot([0, right_x], [0, right_y], 'b--')  # Right vision line\n",
    "\n",
    "# Observations that won't move\n",
    "observation = np.zeros(env.observation_size)\n",
    "observation[0:5] = [10 , 10, 0, 0.1 , 0.1 ]\n",
    "for j in [0, 1, 2, 3, 4, 5]:\n",
    "    base_index = 5  + j * env.num_observed_properties\n",
    "    if env.use_speed_observation:\n",
    "        observation[base_index + 3] = 0                   # speed 1\n",
    "        observation[base_index + 4] = 0                   # speed 2\n",
    "        observation[base_index + 5] = 0                   # prey/predator type\n",
    "    else:\n",
    "        observation[base_index + 3] = 0                   # prey/predator type\n",
    "\n",
    "#  Calculate the x values for each bell curve\n",
    "x_acc = np.linspace(-0.1, env.max_acceleration_prey, 100)\n",
    "x_turn = np.linspace(-env.max_turn, env.max_turn, 100)\n",
    "\n",
    "# Function to update the plot\n",
    "def update_plot_actions(r, theta, heading):\n",
    "    # Create a figure with 2 subplots (one above the other)\n",
    "    fig, ax_actions = plt.subplots(2, 1, figsize=(8, 6))\n",
    "    \n",
    "    for j in [0, 1, 2, 3, 4, 5]:\n",
    "        base_index = 5  + j * env.num_observed_properties\n",
    "        observation[base_index] =    r / env.grid_diagonal   # relative normalized position x\n",
    "        observation[base_index + 1] = theta / (2 * np.pi)    # relative normalized position y\n",
    "        observation[base_index + 2] = heading / (2 * np.pi)  # relative heading (always the same here)\n",
    "\n",
    "    mean_acc, mean_turn, log_std_acc, log_std_turn = policy.compute_single_action(observation)[2][\"action_dist_inputs\"].tolist()\n",
    "    \n",
    "    std_acc = np.exp(log_std_acc)\n",
    "    std_turn = np.exp(log_std_turn)\n",
    "    \n",
    "    # Calculate the y values for each bell curve\n",
    "    y_acc = (1 / (np.sqrt(2 * np.pi) * std_acc)) * np.exp(-0.5 * ((x_acc - mean_acc) / std_acc) ** 2)\n",
    "    y_turn = (1 / (np.sqrt(2 * np.pi) * std_turn)) * np.exp(-0.5 * ((x_turn - mean_turn) / std_turn) ** 2)\n",
    "\n",
    " \n",
    "\n",
    "    # Plotting on the acceleration amplitude subplot\n",
    "    ax_actions[0].plot(x_acc, y_acc, color='blue', linestyle='-', linewidth=2, label='Amplitude')\n",
    "    ax_actions[0].grid(True)\n",
    "    ax_actions[0].legend()\n",
    "    ax_actions[0].set_title(\"action amplitude\")\n",
    "    \n",
    "    # Plotting on the turn angle subplot\n",
    "    ax_actions[1].plot(x_turn, y_turn, color='red', linestyle='--', linewidth=2, label='Orientation')\n",
    "    ax_actions[1].grid(True)\n",
    "    ax_actions[1].legend()\n",
    "    ax_actions[1].set_title(\"action turn\")\n",
    "    \n",
    "    # Adjust layout to not overlap plots\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Interactive sliders for radius and angle\n",
    "radius_slider = FloatSlider(min=0, max=env.max_seeing_distance, step=0.1, value=20, description='Distance (r)', orientation='horizontal')\n",
    "theta_slider = FloatSlider(min=-np.pi, max=np.pi, step=0.1, value= 0.5, description='Angle (θ)', orientation='horizontal')\n",
    "heading_slider = FloatSlider(min=-np.pi, max=np.pi, step=0.1, value= 0.5, description='Heading (θ)', orientation='horizontal')\n",
    "\n",
    "\n",
    "# Create the interactive plot\n",
    "interactive_plot1 = interactive(update_plot_observation, r=radius_slider, theta=theta_slider, heading=heading_slider)\n",
    "interactive_plot2 = interactive(update_plot_actions, r=radius_slider, theta=theta_slider, heading=heading_slider)\n",
    "\n",
    "# Extract the output and set the height\n",
    "output1 = interactive_plot1.children[-1]\n",
    "output2 = interactive_plot2.children[-1]\n",
    "\n",
    "# Adjust the output layout to have the sliders at the bottom of the first plot\n",
    "slider_box = VBox([radius_slider, theta_slider, heading_slider])\n",
    "left_box = VBox([output1, slider_box])\n",
    "full_plot = HBox([left_box, output2])\n",
    "display(full_plot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a2c5fb8-3105-4ae8-b53b-ad494f147b36",
   "metadata": {},
   "source": [
    "# Atraction, aligment and collision zones "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "756b6f59-0866-42ca-b313-c34eb0c49c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "\n",
    "def compute_zones(r, theta):\n",
    "    # Define a range of heading values to simulate, e.g., [0, 2*pi] in steps\n",
    "    heading_values = np.linspace(-np.pi, np.pi, num=10)  # Adjust num for finer granularity\n",
    "    alignment_info = []\n",
    "\n",
    "    for heading in heading_values:\n",
    "        observation = np.zeros(env.observation_size)\n",
    "        # Example initialization, adjust as needed\n",
    "        observation[0:5] = [10 , 10 , 0, 0.1 , 0.1 ]\n",
    "\n",
    "        # Set up observation based on the current simulated heading\n",
    "        for j in range(6):  # Assuming 6 observations for simplification\n",
    "            base_index = 5 + j * env.num_observed_properties\n",
    "            observation[base_index] = r / env.grid_diagonal\n",
    "            observation[base_index + 1] = theta / (2 * np.pi)\n",
    "            observation[base_index + 2] = heading / (2 * np.pi)\n",
    "            if env.use_speed_observation:\n",
    "                observation[base_index + 3:base_index + 6] = [0, 0, 0]\n",
    "            else:\n",
    "                observation[base_index + 3] = 0\n",
    "\n",
    "        action_dist_inputs = policy.compute_single_action(observation)[2][\"action_dist_inputs\"].tolist()\n",
    "        mean_acc, mean_turn, log_std_acc, log_std_turn = action_dist_inputs\n",
    "        std_acc = np.exp(log_std_acc)\n",
    "        std_turn = np.exp(log_std_turn)\n",
    "\n",
    "        aligment_turn = (mean_turn - heading) % (2 * np.pi) - np.pi\n",
    "        is_aligned = (-env.max_turn <= mean_turn <= env.max_turn) and np.sign(mean_turn) == np.sign(aligment_turn)\n",
    "\n",
    "        # Capture alignment information\n",
    "        alignment_info.append(is_aligned)\n",
    "\n",
    "    alignment_ratio = sum(alignment_info) / len(alignment_info)\n",
    "    \n",
    "    # Define thresholds\n",
    "    collision_threshold = std_turn * 2\n",
    "    attraction_threshold = std_turn * 0.5\n",
    "    alignment_threshold = 0.6\n",
    "\n",
    "    # Zone computation\n",
    "    if alignment_ratio >= alignment_threshold:\n",
    "        return 'alignment', alignment_ratio \n",
    "    elif np.abs(mean_turn - theta) <= attraction_threshold:\n",
    "        zone = 'attraction'\n",
    "    elif np.abs(mean_turn - theta) > collision_threshold:\n",
    "        zone = 'collision'\n",
    "    else:\n",
    "        zone = 'indeterminate'  # Conditions for other zones not met\n",
    "\n",
    "    return zone\n",
    "\n",
    "\n",
    "def plot_zones():\n",
    "    plt.figure(figsize=(10, 10))\n",
    "\n",
    "    # Assuming env.max_seeing_distance, env.stage_size, and other parameters are defined\n",
    "    for r in np.linspace(0, env.max_seeing_distance, 50):\n",
    "        for theta in np.linspace(-env.max_seeing_angle, env.max_seeing_angle, 100):\n",
    "            zone = compute_zones(r, theta)\n",
    "            x = r * np.cos(theta)\n",
    "            y = r * np.sin(theta)\n",
    "\n",
    "            if zone == 'collision':\n",
    "                color = 'red'\n",
    "            elif zone == 'attraction':\n",
    "                color = 'green'\n",
    "            elif zone == 'indeterminate':\n",
    "                color = 'grey'\n",
    "            else:  # alignment\n",
    "                color = 'blue'\n",
    "\n",
    "            plt.plot(x, y, marker='.', color=color)\n",
    "\n",
    "    # Plot formatting, assuming you've defined these variables\n",
    "    plt.xlim(-env.stage_size, env.stage_size)\n",
    "    plt.ylim(-env.stage_size, env.stage_size)\n",
    "    plt.axhline(0, color='black', linewidth=0.5)\n",
    "    plt.axvline(0, color='black', linewidth=0.5)\n",
    "\n",
    "    # Add legend\n",
    "    collision_patch = patches.Patch(color='red', label='Collision Zone')\n",
    "    attraction_patch = patches.Patch(color='green', label='Attraction Zone')\n",
    "    alignment_patch = patches.Patch(color='blue', label='Alignment Zone')\n",
    "    plt.legend(handles=[collision_patch, attraction_patch, alignment_patch])\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "# Call the plotting function\n",
    "plot_zones()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51512a01-91fc-459f-9530-311cdd7ea36b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Brouillon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af37094-77be-4a80-b60a-9a1f5fb1d8d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = algo.get_policy(\"prey\")\n",
    "\n",
    "policy.make_model_and_action_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c12a046-654b-498b-a19e-7a7455c6b287",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "fc_net_output = policy.model.encoder.actor_encoder.net.mlp(torch.Tensor(toy_vector))\n",
    "\n",
    "outputs = policy.model.pi.net(fc_net_output).tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b34739d-4ca3-49b0-bbcf-41aa242c1ea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "policy.model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ae2e025-8041-4c6f-81c6-5da0736fc38f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_config = policy.model.config\n",
    "catalog = model_config.catalog_class(model_config.observation_space, model_config.action_space, model_config.model_config_dict)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5efac95b-04ae-4dc1-b08e-b798f55b0127",
   "metadata": {},
   "outputs": [],
   "source": [
    "action_dist_class = catalog.get_action_dist_cls(framework=\"torch\")\n",
    "\n",
    "fc_net_output = policy.model.encoder.actor_encoder.net.mlp(torch.Tensor(toy_vector))\n",
    "\n",
    "action_dist_inputs = policy.model.pi.net(fc_net_output)\n",
    "\n",
    "action_dist = action_dist_class.from_logits(action_dist_inputs)\n",
    "actions = action_dist.sample().numpy()\n",
    "actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dd460f7-7eb2-4472-b037-a6e63cdc4d5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "action_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2362d632-5d61-4144-b5ce-bf6d831bb277",
   "metadata": {},
   "outputs": [],
   "source": [
    "action_dist_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89dea25e-be6b-4cb0-b78c-0749975c65a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "policy.model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97b71db8-253c-421a-b8f8-40b402efe2ca",
   "metadata": {},
   "source": [
    "# Network visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "303536c0-fbd1-4c5e-a5f0-1d2c8e15ae1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example shape\n",
    "policy_id = \"prey\"\n",
    "print(algo.get_policy(policy_id).get_weights()['pi.net.mlp.0.bias'].shape)\n",
    "# We create a subdictionnary with the interresting layers\n",
    "actor_weights = {}\n",
    "for key, value in algo.get_policy(policy_id).get_weights().items():\n",
    "    if \"critic_encoder\" not in key and \"vf.\" not in key:\n",
    "        actor_weights[key] = value\n",
    "\n",
    "actor_weights.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8ceb75c-8281-4826-9931-4c89af24eba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from graph_tool.all import *\n",
    "\n",
    "def create_graph(neural_network):\n",
    "    g = Graph(directed=True)\n",
    "    \n",
    "    # Create property maps for vertex and edge labels and edge width\n",
    "    v_label = g.new_vertex_property(\"string\")\n",
    "    e_width = g.new_edge_property(\"double\")\n",
    "    pos = g.new_vertex_property(\"vector<double>\")\n",
    "    \n",
    "    max_neurons = max(len(neural_network[key]) for key in neural_network if 'weight' in key)\n",
    "\n",
    "    def add_layer_and_set_positions(neurons, x_pos, pos):\n",
    "        layer_vertices = [g.add_vertex() for _ in neurons]\n",
    "        starting_y = (max_neurons - len(layer_vertices)) / 2\n",
    "        for i, v in enumerate(layer_vertices):\n",
    "            pos[v] = (x_pos, starting_y + len(layer_vertices) - 1 - i)\n",
    "        return layer_vertices\n",
    "    \n",
    "    ## VERTEX ##\n",
    "    # Input Layer\n",
    "    layers = [add_layer_and_set_positions(neural_network['encoder.actor_encoder.net.mlp.0.weight'].T, 0, pos)]\n",
    "    \n",
    "    # Hiden Layers\n",
    "    x_gap = 20 # gap between layers\n",
    "    biases_keys = [key for key in neural_network if \".bias\" in key and \"actor_encoder\" in key]\n",
    "    for i, biases_key in enumerate(biases_keys):\n",
    "        # Add vertices for the current layer and set their positions\n",
    "        layers.append(add_layer_and_set_positions(neural_network[biases_key], x_gap*(i+1), pos))\n",
    "        \n",
    "    # Output Layer\n",
    "    output_neurons = add_layer_and_set_positions(neural_network['pi.net.mlp.0.bias'], x_gap*4, pos)\n",
    "\n",
    "    ## EDGES ##\n",
    "    # Set labels and add edges for input-hidden and hidden-hidden layer\n",
    "    weights_keys = [key for key in neural_network if \".weight\" in key and \"actor_encoder\" in key]\n",
    "    for k, weights_key in enumerate(weights_keys):\n",
    "        for i, hidden_neuron in enumerate(layers[k]):\n",
    "            for j, next_hidden_neuron in enumerate(layers[k+1]):\n",
    "                e = g.add_edge(hidden_neuron, next_hidden_neuron)\n",
    "                weight = neural_network[weights_key].T[i][j]\n",
    "                e_width[e] = weight\n",
    "\n",
    "    # Set labels and add edges for hidden-output layer\n",
    "    for j, output_neuron in enumerate(output_neurons):\n",
    "        for i, hidden_neuron in enumerate(layers[-1]):\n",
    "            e = g.add_edge(hidden_neuron, output_neuron)\n",
    "            weight = neural_network['pi.net.mlp.0.weight'].T[i][j]\n",
    "            e_width[e] = weight\n",
    "\n",
    "    \n",
    "    # LABELS\n",
    "    # Set neuron labels (optional, for clarity)\n",
    "    for v in layers[0]:\n",
    "        v_label[v] = \"I\"\n",
    "    for k, hidden_neurons in enumerate(layers):\n",
    "        for v in hidden_neurons:\n",
    "            v_label[v] = \"H\"\n",
    "    for v in output_neurons:\n",
    "        v_label[v] = \"O\"\n",
    "\n",
    "    return g, pos, v_label, e_width\n",
    "\n",
    "# Example usage with the same nn_wandb\n",
    "g, pos, v_label, e_width = create_graph(actor_weights)\n",
    "\n",
    "# Draw the graph\n",
    "graph_draw(g, pos=pos, vertex_text=v_label, edge_text=None, edge_pen_width=e_width, vertex_size=15, vertex_font_size=10, edge_font_size=10, output_size=(800, 800))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bcd1d81-a0de-428c-93a0-1d0797c1f64e",
   "metadata": {},
   "source": [
    "## Can we reproduce the action by doing a forward pass on the network ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86f0920d-f9f5-4a2b-9f6b-982b8e640289",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def neural_net_activation(input_vector, neural_dic):\n",
    "    # Ensuring the input is a numpy array\n",
    "    input_vector = np.array(input_vector)\n",
    "\n",
    "    # Layer activations\n",
    "    activations = {}\n",
    "\n",
    "    # First three hidden layers\n",
    "    for i in range(3):\n",
    "        weight_key = f\"encoder.actor_encoder.net.mlp.{2*i}.weight\"\n",
    "        bias_key = f\"encoder.actor_encoder.net.mlp.{2*i}.bias\"\n",
    "\n",
    "        if i == 0:\n",
    "            layer_input = input_vector\n",
    "        else:\n",
    "            layer_input = activations[f\"layer_{i}\"]\n",
    "\n",
    "        # Calculate the layer output\n",
    "        z = np.dot(layer_input, neural_dic[weight_key].T) + neural_dic[bias_key]\n",
    "        activations[f\"layer_{i+1}\"] = np.tanh(z)\n",
    "\n",
    "    # Output layer\n",
    "    weight_key = \"pi.net.mlp.0.weight\"\n",
    "    bias_key = \"pi.net.mlp.0.bias\"\n",
    "    output_layer_input = activations[\"layer_3\"]\n",
    "    z = np.dot(output_layer_input, neural_dic[weight_key].T) + neural_dic[bias_key]\n",
    "    activations[\"output\"] = np.tanh(z)\n",
    "\n",
    "    return activations\n",
    "\n",
    "toy_vector = [0.95600354, 0.36926809,\n",
    " 0.07873756, 0.38905384, 0.92927526, 0.08713003, 0.61284082,\n",
    " 0.82746801, 0.36660529, 0.89503505, 0.04707359, 0.97744959,\n",
    " 0.28597701, 0.96454964, 0.29067754, 0.74090134, 0.8245886, 0.96454964, 0.29067754, 0.74090134, 0.8245886,\n",
    "             0.96454964, 0.29067754, 0.74090134, 0.8245886, 0.96454964, 0.29067754, 0.29067754, 0.29067754]\n",
    "\n",
    "activations = neural_net_activation(toy_vector, actor_weights)\n",
    "\n",
    "activations[\"output\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a78f2e9-a06e-4051-a071-ac272e1005c7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a9ba9b1-3d6f-416f-9aff-3fc29b6ef6aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}