{"cells":[{"cell_type":"code","execution_count":97,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":62168,"status":"ok","timestamp":1702900313432,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"LFhdskOWb-HT","outputId":"cb624f4f-3eb6-4ae2-9c7d-7edda3b28b5d"},"outputs":[],"source":["#@title 1. INIT - Imports and clone data from GitHub repo\n","\n","import gc\n","import math as mt\n","import scipy as sp\n","import random\n","import numpy as np\n","import matplotlib.pyplot as plt\n","from matplotlib.font_manager import FontProperties\n","from collections import namedtuple, deque\n","import networkx as nx\n","import pickle as pkl\n","from sklearn.decomposition import PCA\n","from sklearn.manifold import TSNE\n","from sklearn.preprocessing import StandardScaler\n","from pycolormap_2d import ColorMap2DZiegler\n","import os #for creating directories\n","from scipy.optimize import curve_fit\n","from typing import Union, Optional\n","import umap.umap_ as umap\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torch.utils.data import DataLoader, Dataset\n","#this is set for the printing of Q-matrices via console\n","torch.set_printoptions(precision=3, sci_mode=False, linewidth=100)\n","\n","from tqdm.auto import tqdm, trange\n","\n","\n","#Uncomment if GPU is to be used - right now use CPU, as we have very small networks and for them, CPU is actually faster\n","device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","\n","#suppress scientific notation in printouts\n","np.set_printoptions(suppress=True)\n","\n","#location of the data folder w.r.t. the notebook\n","file_loc: str=\"/home/tobi/codebase/language_emergence/multi_agent_language/data/\""]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":288,"status":"ok","timestamp":1702900606217,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"ngt2rljKokAA","outputId":"62ea7ac5-4b79-4046-e5ad-59e57e12d5a8"},"outputs":[],"source":["#@title 2. INIT - Changeable parameters\n","\n","#generate a new \"language\" with the autoencoder, or load it from a file\n","language_code: str=\"nonlinear_goallocs0_zeta5_language2\" #code to get autoencoder parameters from a file\n","\n","#closing the loop\n","language_code_closingloop=\"nonlinear_goallocs0_zeta5_language0\" #language code for which language we should plot the results of \"closing the loop\"\n","\n","#autoencoder loss plots\n","zeta_lossplot=5 #change the hyperparameter zeta here to see the loss plots for different values (1,2,5,10)\n","\n","#plot the student performances for different groups of goal locations trained\n","folders_goalloc_plots=[f\"nonlinear_goallocs{i}_zeta5_factor1\" for i in range(7)] #the goal groups 0-6 have the same order as in the plots in the paper\n","language_nr_goalloc_plots=5 #how many languages per group of goal locations?\n","chuck_out=False #chuck out languages in which the performance of the informed student is below the misinformed student or below the random walker?\n","\n","#read Q-matrices from a file\n","qmat_read_code: str=\"training_4x4\" #either training4x4 or test4x4\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":2,"status":"ok","timestamp":1702900606907,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"BEuD0mnvGgu6"},"outputs":[],"source":["#@title 3. INIT - Fixed parameters\n","\n","grid_dim: int=4 #side length of the square gridworld\n","n_actions: int=4 #how many actions are possible in each state\n","student_init: int=0 #initial state where the student is always starting - states are indexed as follows (example for the 4x4 mazes):\n","# 12 13 14 15\n","# 8 9 10 11\n","# 4 5 6 7\n","# 0 1 2 3\n","lava: bool=False #do we use lava states - i.e. accessible wall states - (True) or wall states (False)? (we now always use \"False\")\n","data_shape: 'tuple[int,int,int]'=(n_actions,grid_dim,grid_dim) #Q matrix shape\n","\n","\n","#rewards in the gridworld\n","step_reward: float=-0.1 #for taking a step\n","goal_reward: float=2. #for reaching the goal\n","wall_reward: float=-0.5 #for bumping into a wall\n","\n","\n","#teacher network parameters for learning Q-matrices\n","qmat_gen: bool=False #generate new Q-matrices?\n","accuracy=20 #how accurate do the final Q-matrices have to be compared to the \"perfect\" ones, which can be calculated by table lookup? (in terms of vector 1-norm)\n","max_attempts=3 #limit the attempts we give the teacher to calculate a Q-matrix that lies within the accuracy bound\n","wall_state_dict={0:[]} #only in case of new Q-matrix generation, we index all the mazes, and calculate a Q-matrix for every possible goal location in each maze\n","gamma_bellman: float=0.99 #temporal discount factor used in the Bellman equation of the teacher networks\n","L: int=50 #short term memory size of teacher learning Q-matrices\n","lr_teacher: float=3e-3 #learning rate for the teachers\n","\n","\n","#autoencoder network parameters for generating the language\n","language_gen: bool=False #generate a new language?\n","K: int=5 #length of the message vectors\n","gamma_sparse=1/20*mt.sqrt(grid_dim**2*n_actions/K) #how important is the sparsity loss compared to reconstruction loss\n"," #we have the square root here to compare individual entries (in both cases take the overall 2-norm)\n","kappa=1/500 #balances how much focus is put on regularization, compared to the probability of finding the goal (for the student)\n","learning_rate_autoenc=5e-4 #learning rate of the autoencoder-student network\n","training_epochs=50 #number of training epochs for the autoencoder\n","zeta_std=5 #the zeta parameter (std: \"student\") for the loss function\n","\n","#message space PCA (and t-SNE) plots\n","do_tsne_message_plots=True #do t-SNE plots as well?\n","do_umap_message_plots=True #do UMAP plots as well?\n","plot_worlds=range(16) #all these worlds are included in the main plot - can not be more than 20 at the moment due to the color cycler -> range(16) are all training worlds\n","plot_worlds_single=[6] #create individual plots for all of the worlds included in this list\n","nonlinear_ae_plots=True if language_code.__contains__(\"nonlinear\") else False #does the autoencoder have nonlinear activations?\n","nonlinear_std_plots=True if (language_code.__contains__(\"nonlinear\") or language_code.__contains__(\"linear_ae\")) else False #does the student have nonlinear activations?\n","save_message_plots=False #save the plots?\n","\n","#evaluate the student agent's performance (calculate goal finding success rates)\n","student_evaluate: bool=False\n","save_rates: bool=True #save the goal finding success rates\n","\n","#autoencoder loss plots\n","epskip=50 #skip first few epochs to have narrower loss range and see details better\n","save_autoenc_lossplots=False #save the autoencoder plots?\n","\n","#plot the student performances for different amounts of steps allowed for the students (\"stepfactors\")\n","folder_stepfactor_plots=\"nonlinear_goallocs0_zeta5_factor1\"\n","language_nr_stepfactor_plots=5\n","stepfactor_list_plots=[1,1.5,2,2.5,3,3.5,4]\n","rdrates_list_plots=[]\n","save_stepfactor_plots=False #save the plots?\n","\n","#plot the student performances for different groups of goal locations trained\n","goal_groups_plots=np.array([0,1,2,3,4,5,6]) #include those groups in the plot\n","stepfactor_goalloc_plots=2\n","rdrate_goalloc_plots=0.25\n","save_goalloc_plots=False #save the plots?\n","\n","#closing the loop\n","closingloop_nonlinear_ae=True #do the neurons in the autoencoder have nonlinear activations?\n","closingloop_nonlinear_std=True #do the neurons in the student have nonlinear activations?\n","save_closingloop_plots=False #save the plots?\n","\n","#topographic similarity analysis (Fig. 3b,c)\n","norm_topo=2 #which norm to use for meaning distances, i.e. task vectors and Q-matrices\n","n_bins_topo=20 #how many bins on the x-axis of the plot?\n","\n","#entropy analysis through binning (Fig. 3e)\n","H_binlist=np.linspace(8,62,10).astype(int) #list of number of bins to calculate entropy from"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":259,"status":"ok","timestamp":1702900607539,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"ss4KO1ADTbzd"},"outputs":[],"source":["#@title 4. FUNCTIONS - Helper functions\n","\n","def set_seed(seed: Optional[int]=None, seed_torch: bool=True):\n"," {'''\n"," Function that controls randomness by setting a random seed\n"," ---\n"," INPUT:\n"," seed: the random state\n"," seed_torch: If `True` sets the random seed for pytorch tensors, so pytorch module\n"," must be imported\n"," ---\n"," OUTPUT:\n"," nothing, write into file\n"," '''}\n"," if seed is None:\n"," seed = np.random.choice(2 ** 32)\n"," random.seed(seed)\n"," np.random.seed(seed)\n"," if seed_torch:\n"," torch.manual_seed(seed)\n"," torch.cuda.manual_seed_all(seed)\n"," torch.cuda.manual_seed(seed)\n"," torch.backends.cudnn.benchmark = False\n"," torch.backends.cudnn.deterministic = True\n","\n"," print(f'Random seed {seed} has been set.')\n","\n","\n","def write_dict_into_pkl(dictionary: dict,pkl_file: str):\n"," {'''\n"," Write a dictionary into a .txt file\n"," ---\n"," INPUT:\n"," dictionary: the dictionary\n"," txt_file: location of the .txt file we are writing into\n"," ---\n"," OUTPUT:\n"," nothing, write into file\n"," '''}\n"," with open(pkl_file,\"wb\") as f:\n"," pkl.dump(dictionary, f)\n"," f.close()\n","\n","\n","def read_dict_from_pkl(pkl_file: str)->dict:\n"," {'''\n"," Read a dictionary from a .txt file\n"," ---\n"," INPUT:\n"," txt_file: location of the .txt file we are reading from\n"," ---\n"," OUTPUT:\n"," the dictionary we get from the file\n"," '''}\n"," pickle_off = open(pkl_file, 'rb')\n"," dictionary = pkl.load(pickle_off)\n"," pickle_off.close()\n"," return dictionary\n","\n","\n","def moving_average(a: Union[list, np.ndarray], n: int) ->np.ndarray:\n"," {'''\n"," computes the moving average of an array\n"," ---\n"," INPUT\n"," a - the array\n"," n - width of the moving window in steps\n"," ---\n"," OUTPUT\n"," the moving average array\n"," '''}\n"," ret = np.cumsum(a, dtype=float)\n"," ret[n:] = ret[n:] - ret[:-n]\n"," return ret[n - 1:] / n\n","\n","def linfunc(x: Union[float,int],m: Union[float,int],b: Union[float,int])->Union[float,int]:\n"," '''\n"," A simple linear function\n"," ---\n"," INPUT\n"," x - input x-value\n"," m - linear function slope parameter\n"," b - linear function offset parameter\n"," ---\n"," OUTPUT\n"," y=m*x+b\n"," '''\n"," return m*x+b\n","\n","def tanhfunc(x: Union[float,int],a: Union[float,int], c: Union[float,int], b: Union[float,int])->Union[float,int]:\n"," '''\n"," Hyperbolic tangent function with scaling parameters\n"," ---\n"," INPUT\n"," x - input x-value\n"," a - parameter\n"," b - scaling parameter (y-axis)\n"," ---\n"," OUTPUT\n"," y=tanh(a*x)\n"," '''\n"," return b*np.tanh(a*(x-1)**c)\n","\n","def logfunc(x: Union[float,int],a: Union[float,int], c: Union[float,int], b: Union[float,int])->Union[float,int]:\n"," '''\n"," Logistic function with scaling parameters\n"," ---\n"," INPUT\n"," x - input x-value\n"," a - parameter\n"," b - scaling parameter (y-axis)\n"," ---\n"," OUTPUT\n"," y=(2*b)/(1+e^(-a*x))-b\n"," '''\n"," return 2*b/(1+np.exp(-a*(x-1)**c))-b\n","\n","def hypfunc(x: Union[float,int],a: Union[float,int], c: Union[float,int], b: Union[float,int])->Union[float,int]:\n"," '''\n"," Hyperbola with scaling parameter\n"," ---\n"," INPUT\n"," x - input x-value\n"," a - parameter\n"," c - parameter\n"," b - scaling parameter (y-axis)\n"," ---\n"," OUTPUT\n"," y=1/(a*x)+b\n"," '''\n"," return -1/(a*(x-1)**c)+b\n","\n","def arctanfunc(x: Union[float,int],a: Union[float,int], c: Union[float,int], b: Union[float,int])->Union[float,int]:\n"," '''\n"," Hyperbola with scaling parameter\n"," ---\n"," INPUT\n"," x - input x-value\n"," a - parameter\n"," c - parameter\n"," b - scaling parameter (y-axis)\n"," ---\n"," OUTPUT\n"," y=1/(a*x)+b\n"," '''\n"," return (2/mt.pi)*b*np.arctan(a*(x-1)**c)\n","\n","def constfunc(x, a):\n"," return 0*x+a\n","\n","def flattening_spotter(data: Union[list, np.ndarray], fit_size: int, certainty: float) ->int:\n"," '''\n"," This function finds the point when a series of data points (data) has converged/levelled/flattened\n"," It assumes that by the end of the data series (minus fit_size data points), convergence has definitely happened!!\n"," ---\n"," INPUT\n"," data - list/array of numbers. the data series\n"," fit_size - the number of data points taken into account for each linear fit\n"," certainty - if the slope plus this number of standard deviations is below zero, we stop and declare that the data is no longer level\n"," ---\n"," OUTPUT\n"," learned_index - the first index where the data series has become level\n"," '''\n","\n"," learned_index: int=0\n"," for l in range(len(data)-fit_size,-1,-1):\n"," popt, pcov=curve_fit(linfunc,range(l,l+fit_size,1),data[l:l+fit_size]) #fit linear function\n"," if popt[0]<1e-3 and popt[0]+certainty*mt.sqrt(pcov[0,0])<1e-3: #if the slope is negative with given certainty, we stop\n"," learned_index=l+1 #+1 step, don't include the data that is not level anymore\n"," break\n"," return learned_index\n","\n","\n","def to_interval(x: float,left: float,right: float)->float:\n"," '''\n"," Fits data point into an interval (left outliers to left boundary, right outliers to right boundary)\n"," ---\n"," INPUT\n"," x - the data point\n"," left - left interval boundary\n"," right - right interval boundary\n"," ---\n"," OUTPUT\n"," y - the position of x in the interval\n"," '''\n","\n"," return min(right,max(x,left))\n","\n","\n","def set_axes_equal(ax):\n"," '''Make axes of 3D plot have equal scale so that spheres appear as spheres,\n"," cubes as cubes, etc.. This is one possible solution to Matplotlib's\n"," ax.set_aspect('equal') and ax.axis('equal') not working for 3D.\n","\n"," Input\n"," ax: a matplotlib axis, e.g., as output from plt.gca().\n"," '''\n","\n"," x_limits = ax.get_xlim3d()\n"," y_limits = ax.get_ylim3d()\n"," z_limits = ax.get_zlim3d()\n","\n"," x_range = abs(x_limits[1] - x_limits[0])\n"," x_middle = np.mean(x_limits)\n"," y_range = abs(y_limits[1] - y_limits[0])\n"," y_middle = np.mean(y_limits)\n"," z_range = abs(z_limits[1] - z_limits[0])\n"," z_middle = np.mean(z_limits)\n","\n"," # The plot bounding box is a sphere in the sense of the infinity\n"," # norm, hence I call half the max range the plot radius.\n"," plot_radius = 0.5*max([x_range, y_range, z_range])\n","\n"," ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])\n"," ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])\n"," ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])\n","\n","def get_message_list(conv_autoenc, language_code, q_matrix_dict):\n"," #create the message list corresponding to a certain language\n"," #load stored autoencoder network parameters\n"," autoencoder = conv_autoenc(data_shape, K, nonlinear_ae_plots, nonlinear_std_plots).to(device)\n"," autoencoder.load_state_dict(torch.load(file_loc+\"autoencoder/autoencoder network parameters/\"+f\"params_autoenc{language_code}.pt\", weights_only=True))\n"," autoencoder.eval()\n"," #create a message dictionary, with indices corresponding to the task indices\n"," message_dict={}\n"," for task_index,q_matrix in q_matrix_dict.items():\n"," q_matrix=torch.unsqueeze(q_matrix,0) #need this because the autoencoder always expects batches of inputs!\n"," message=autoencoder.encode(q_matrix)[0]\n"," message_dict[task_index]=torch.flatten(message).detach().numpy()\n","\n"," return list(message_dict.values()), autoencoder\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":3,"status":"ok","timestamp":1702900608467,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"UcspfZ7eih-Z"},"outputs":[],"source":["#@title 5. FUNCTIONS - Gridworld class and state representation transformations\n","\n","class SquareGridworld():\n"," \"\"\"\n"," the class for the gridworlds\n","\n"," n**2 states (n-by-n gridworld) -> n is the grid dimension\n","\n"," The mapping from state to the grid is as follows:\n"," n(n-1) n(n-1)+1 ... n^2-1\n"," ... ... ... ...\n"," n n+1 ... 2n-1\n"," 0 1 ... n-1\n","\n"," Actions 0, 1, 2, 3 correspond to right, up, left, down (always exactly one step)\n","\n"," -Agent starts at the init_state\n"," -Landing in the goal_state gets a reward of goal_reward and ends the episode\n"," -Bumping into wall states or the map border incurs a reward of wall_reward\n"," -In case of lava=False, the agent bounces back from walls, while in case of lava=True wall states are accessible (the outside boundary can never be crossed)\n"," -Each step additionally incurs a reward of step_reward\n"," \"\"\"\n"," def __init__(self,init_state: int,goal_state: int,wall_states: 'list[int]', lava: bool):\n","\n"," self.init_state: int=init_state\n"," self.goal_state: int=goal_state\n"," self.wall_states: 'list[int]'=wall_states\n"," self.n_states: int = grid_dim**2\n"," self.lava: bool = lava\n","\n","\n"," def get_outcome(self, state: int, action: int)-> 'tuple[Optional[int], float]':\n"," '''\n"," given a state and an action, this returns the next state and the immediate reward for the action\n"," ---\n"," INPUT\n"," state: the state the agent is in\n"," action: the action taken\n"," ---\n"," OUTPUT\n"," next_state - the next state\n"," reward - the reward for the action taken\n"," '''\n","\n"," #if the goal is reached, we get the goal_reward and the episode ends\n"," if state == self.goal_state:\n"," reward: float = goal_reward\n"," next_state = None #terminates the episode\n"," return next_state, reward\n","\n"," #get the next state before taking into account walls or outside boundary\n"," next_state_dict={0:state+1, 1:state+grid_dim, 2:state-1, 3:state-grid_dim}\n"," #for all actions, this dictionary stores a boolean indicating whether we cross the map border executing this action in our current state\n"," cross_boundary_dict={0:state % grid_dim == grid_dim-1, 1:state >= grid_dim*(grid_dim-1), 2:state % grid_dim ==0, 3:state 'dict[ tuple[int,float] , tuple[Optional[int], float] ]':\n"," '''\n"," returns a dictionary where for every possible combination of state and action we get the next state and\n"," corresponding immediate reward\n"," ---\n"," OUTPUT\n"," outcomes - the dictionary keyed by state-action combo, whose values are the next states and rewards\n"," '''\n"," outcomes = {(s, a): self.get_outcome(s,a) for s in range(self.n_states) for a in range(n_actions)}\n"," return outcomes\n","\n","\n","\n","def state_int_to_tuple(state_int: Optional[int]) -> Optional[torch.tensor]:\n"," '''\n"," Gets state tuple representation (coordinates) from integer representation\n"," ---\n"," INPUT\n"," state_int - the state integer representation\n"," ---\n"," OUTPUT\n"," state - the state tuple representation\n"," '''\n"," if state_int==None:\n"," return None\n"," else:\n"," cval: float=(grid_dim-1)/2 #center (0,0) is in the middle of the grid\n"," sx,sy=state_int%grid_dim-cval, mt.floor(state_int/grid_dim)-cval\n"," state = torch.tensor([[sx,sy]],device=device)\n"," return state\n","\n","def state_tuple_to_int(state: Optional[torch.tensor]) -> Optional[int]:\n"," '''\n"," Gets state integer representation from tuple representation (coordinates)\n"," ---\n"," INPUT\n"," state - the state tuple representation\n"," ---\n"," OUTPUT\n"," state_int - the state integer representation\n"," '''\n"," if state==None:\n"," return None\n"," else:\n"," sx,sy=state\n"," cval: float=(grid_dim-1)/2\n"," sx,sy=sx+cval,sy+cval\n"," state_int=int(sx.item()+grid_dim*sy.item())\n"," return state_int\n","\n","def get_state_tensors(m_len: int) -> torch.tensor:\n"," '''\n"," Concatenate tensors representing particular states in the gridworld so that fewer network forwards are needed\n"," ---\n"," INPUT\n"," m_len - length of the message batch\n"," ---\n"," OUTPUT\n"," state_tensors - the concatenated state tensors\n"," '''\n"," state_tensors=torch.zeros(size=(grid_dim,grid_dim*m_len,2))\n"," for j in range(grid_dim):\n"," for i in range(grid_dim):\n"," s=i*grid_dim+j\n"," s_tup=state_int_to_tuple(s)[0]\n"," for b in range(m_len):\n"," state_tensors[j,i*m_len+b]=s_tup\n","\n"," return state_tensors\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":390,"status":"ok","timestamp":1702900609253,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"CFSH3Lp-i4_5"},"outputs":[],"source":["#@title 6. FUNCTIONS - Transition Memory class\n","\n","\n","\n","Transition = namedtuple('Transition',('state', 'action', 'next_state', 'reward'))\n","\n","class ReplayMemory(object):\n"," '''\n"," a class to store and replay previous transitions of states and actions done by the network\n"," like a memory\n"," '''\n","\n"," def __init__(self, capacity: int):\n"," #Have short- and long-term memory to optimize exploration results\n"," self.memory = deque([],maxlen=capacity)\n","\n"," def push(self, *args):\n"," \"\"\"Save a transition\"\"\"\n"," self.memory.append(Transition(*args))\n","\n"," def sample(self, batch_size: int)->'list[Transition]':\n"," return random.sample(self.memory, batch_size)\n","\n"," def __len__(self) -> int:\n"," return len(self.memory)"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":2,"status":"ok","timestamp":1702900609580,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"MtBd-1QMi88h"},"outputs":[],"source":["#@title 7. FUNCTIONS - DQN class\n","\n","class DQN(nn.Module):\n"," '''\n"," Deep Q Network class\n"," '''\n","\n"," def __init__(self, K: int, zero_init: bool = False):\n"," '''\n"," INPUT\n"," K: length of input message (zero for the teacher)\n"," zero_init: initialize all weights to zero?\n"," '''\n"," super(DQN, self).__init__()\n"," self.lin1 = nn.Linear(2+K, 10, device=device) #input size is 2(gridworld coordinates)+length of message\n"," self.lin2=nn.Linear(10,20, device=device)\n"," self.lin3 = nn.Linear(20,20, device=device)\n"," self.lin4 = nn.Linear(20,n_actions, device=device)\n","\n"," #Initialize all weights (and biases) to zero\n"," if zero_init:\n"," torch.nn.init.zeros_(self.lin1.weight)\n"," torch.nn.init.zeros_(self.lin1.bias)\n"," torch.nn.init.zeros_(self.lin2.weight)\n"," torch.nn.init.zeros_(self.lin2.bias)\n"," torch.nn.init.zeros_(self.lin3.weight)\n"," torch.nn.init.zeros_(self.lin3.bias)\n","\n","\n"," def forward(self, x:torch.tensor)->torch.tensor:\n"," '''\n"," forward pass of the network\n"," ---\n"," INPUT\n"," x - network input, i.e. combination of state and potentially message\n"," ---\n"," OUTPUT\n"," y - the four Q-values, i.e. torch.tensor([Q(s,0),Q(s,1),Q(s,2),Q(s,3)])\n"," '''\n"," x = F.relu(self.lin1(x))\n"," x = F.relu(self.lin2(x))\n"," x = F.relu(self.lin3(x))\n"," x = self.lin4(x)\n"," return x\n","\n","\n","class BiasLayer(nn.Module):\n"," '''\n"," Bias Layer (add bias to individual network nodes/filter positions)\n"," '''\n"," def __init__(self, shape: tuple):\n"," '''\n"," Initialise parameters of bias layer\n"," ---\n"," INPUT\n"," shape: Requisite shape of bias layer\n"," '''\n"," super(BiasLayer, self).__init__()\n"," init_bias = torch.zeros(shape, device=device)\n"," self.bias = nn.Parameter(init_bias, requires_grad=True)\n","\n"," def forward(self, x: torch.tensor)->torch.tensor:\n"," '''\n"," Forward pass\n"," ---\n"," INPUT\n"," x: Input features\n"," ---\n"," OUTPUT\n"," y: Output of bias layer\n"," '''\n"," y=x+self.bias\n"," return y\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":3,"status":"ok","timestamp":1702900611072,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"-rR28ndQjAOl"},"outputs":[],"source":["#@title 8. FUNCTIONS - Policies\n","def select_action_epsgreedy(agent: DQN,state:torch.tensor,message:torch.tensor,epsilon:float)->int:\n"," '''\n"," epsilon-greedy algorithm for action selection\n"," ---\n"," INPUT\n"," agent - the deep-Q-network representing the agent exploring the world\n"," state - the state the agent is currently in (in coordinate representation)\n"," message - the message the agent has received\n"," epsilon - epsilon from the epsilon greedy algorithm\n"," ---\n"," OUTPUT\n"," The action to be taken (0:right, 1:up, 2:left, 3:down)\n"," '''\n"," if random.random() > epsilon:\n"," with torch.no_grad():\n"," input=torch.cat((state[0],message),0)\n"," return agent(input).argmax().item() #do forward pass and return optimal action found\n"," else:\n"," return random.choice(range(4)) #return random action\n","\n","\n","\n","def select_action_optimism(agent: DQN,state_int:int,message:torch.tensor,alpha:float,sa_counts:'dict[tuple[int, int],int]'):\n"," '''\n"," optimism-in-face-of-uncertainty algorithm for action selection\n"," using the unchanged Q-values and no probabilities\n"," ---\n"," INPUT\n"," agent - the deep-Q-network representing the agent exploring the world\n"," state_int - the state the agent is currently in (in integer representation)\n"," message - the message the agent has received\n"," alpha - constant for optimism in face of uncertainty algorithm\n"," sa_counts - counts how often each state-action combination has already been seen by the agent\n"," ---\n"," OUTPUT\n"," The action to be taken (0:right, 1:up, 2:left, 3:down)\n"," '''\n"," state=state_int_to_tuple(state_int)\n"," current_sa_counts=torch.tensor([sa_counts[(state_int, a)] for a in range(4)]).to(device) #pick out counts of current state\n"," with torch.no_grad():\n"," input=torch.cat((state[0],message),0)\n"," qvals=agent(input) #do forward pass\n"," qvals=qvals+alpha/torch.sqrt(current_sa_counts) #add current uncertainties\n"," return qvals.argmax().item() #return optimal action given uncertainties\n","\n","\n","\n","def select_action_eps_trust(agent: DQN, state_int:int, probas_message_matrix:torch.tensor, softy:nn.Softmax, iota:float,\n"," trust:float, current_ep:int, total_eps:int)->'tuple[torch.tensor, torch.tensor, int]':\n"," '''\n"," global trust parameter for teacher and epsilon greedy exploration\n"," ---\n"," INPUT\n"," agent - the deep-Q-network representing the agent exploring the world\n"," state_int - the state the agent is currently in (in integer representation)\n"," probas_message_matrix - matrix of action probabilities according to the message\n"," softy - a softmax function\n"," iota - parameter to adjust student Q-values, so that softmax gives reasonable probabilities\n"," trust - tells how much the student trusts the teacher (between 0 and 1)\n"," current_ep - current episode number\n"," total_eps - total episode number (both parameters relevant for decaying epsilon in epsilon-greedy)\n"," ---\n"," OUTPUT\n"," probas_message - action probabilities in current state according to the message\n"," probas_student - action probabilities in current state according to the student's own learning\n"," action - next action that will be executed by the student\n"," '''\n"," state=state_int_to_tuple(state_int)\n"," with torch.no_grad():\n"," #1.message\n"," probas_message=probas_message_matrix[:,mt.floor(state_int/grid_dim),state_int%grid_dim] #filter out current state\n"," #2.student\n"," qvals_student=agent(state[0])\n"," qvals_student-=min(qvals_student)*torch.ones(n_actions).to(device)\n"," probas_student=softy(iota*qvals_student) #multiply by constant iota for suitable probabilities from Q-values\n"," #3.combination\n"," probas=trust*probas_message+(1-trust)*probas_student\n"," #4.randomness\n"," r=random.random()\n"," if r<0.4*(1-current_ep/total_eps)+0.4*(current_ep/total_eps): #decaying eps-greedy according to episode number\n"," action=random.choice([0,1,2,3])\n"," else:\n"," action=torch.argmax(probas).item()\n"," return probas_message, probas_student, action\n","\n","\n","\n","def select_action_eps_trust_uninfo(agent: DQN, state_int:int, softy:nn.Softmax, iota:float, current_ep:int, total_eps:int)->'tuple[torch.tensor,int]':\n"," '''\n"," mirror image of above function, just without a message (i.e. without a teacher)\n"," ---\n"," INPUT\n"," agent - the deep-Q-network representing the agent exploring the world\n"," state_int - the state the agent is currently in (in integer representation)\n"," softy - a softmax function\n"," iota - parameter to adjust student Q-values, so that softmax gives reasonable probabilities\n"," current_ep - current episode\n"," total_eps - total episode number (both parameters relevant for decaying epsilon in epsilon-greedy)\n"," ---\n"," OUTPUT\n"," probas_student - action probabilities in current state according to the student's own learning\n"," action - next action that will be executed by the student\n"," '''\n"," state=state_int_to_tuple(state_int)\n"," with torch.no_grad():\n"," #1.student\n"," qvals_student=agent(state[0])\n"," qvals_student-=min(qvals_student)*torch.ones(n_actions).to(device)\n"," probas=softy(iota*qvals_student) #multiply by constant iota for suitable probabilities from Q-values\n"," #2.randomness\n"," r=random.random()\n"," if r<0.4*(1-current_ep/total_eps)+0.4*(current_ep/total_eps): #decaying epsilon greedy according to episode number\n"," action=random.choice([0,1,2,3])\n"," else:\n"," action=torch.argmax(probas).item()\n"," return probas, action\n","\n","\n","\n","def select_action_opt_trust(agent: DQN, state_int:int, probas_message_matrix: torch.tensor, softy:nn.Softmax, iota:float,\n"," trust: float, alpha: float, sa_counts:'dict[tuple[int, int],int]')->'tuple[torch.tensor, torch.tensor, int]':\n"," '''\n"," global trust parameter for teacher and optimism in face of uncertainty exploration\n"," ---\n"," INPUT\n"," agent - the deep-Q-network representing the agent exploring the world\n"," state_int - the state the agent is currently in (in integer representation)\n"," probas_message_matrix - matrix of action probabilities according to the message\n"," softy - a softmax function\n"," iota - parameter to adjust student Q-values so that softmax gives reasonable probabilities\n"," trust - tells how much the student trusts the teacher (between 0 and 1)\n"," alpha - constant for optimism in face of uncertainty algorithm\n"," sa_counts - counts how often each state-action combination has already been seen by the agent\n"," ---\n"," OUTPUT\n"," probas_message - action probabilities in current state according to the message\n"," probas_student - action probabilities in current state according to the student's own learning\n"," action - next action that will be executed by the student\n"," '''\n"," state=state_int_to_tuple(state_int)\n"," current_sa_counts=torch.tensor([sa_counts[(state_int, a)] for a in range(4)]).to(device) #pick out counts of current state\n"," with torch.no_grad():\n"," #1.message\n"," probas_message=probas_message_matrix[:,mt.floor(state_int/grid_dim),state_int%grid_dim] #filter out current state\n"," #2.student\n"," qvals_student=agent(state[0])\n"," qvals_student-=min(qvals_student)*torch.ones(n_actions).to(device)\n"," probas_student=softy(iota*qvals_student) #multiply by constant iota for suitable probabilities from Q-values\n"," #3.combination including optimism in the face of uncertainty\n"," probas=trust*probas_message+(1-trust)*(probas_student+alpha/(current_sa_counts+1))\n"," #argmax\n"," action=torch.argmax(probas).item()\n"," return probas_message, probas_student, action\n","\n","\n","\n","def select_action_opt_trust_uninfo(agent: DQN, state_int:int, softy:nn.Softmax, iota:float, alpha:float, sa_counts:'dict[tuple[int, int],int]')->int:\n"," '''\n"," mirror image of above function, just without a message (i.e. without a teacher)\n"," ---\n"," INPUT\n"," agent - the deep-Q-network representing the agent exploring the world\n"," state_int - the state the agent is currently in (in integer representation)\n"," softy - a softmax function\n"," iota - parameter to adjust student Q-values so that softmax gives reasonable probabilities\n"," alpha - constant for optimism in face of uncertainty algorithm\n"," sa_counts - counts how often each state-action combination has already been seen by the agent\n"," ---\n"," OUTPUT\n"," probas - action probabilities in state s according to the student's own learning\n"," action - next action that will be executed by the student\n"," '''\n"," state=state_int_to_tuple(state_int)\n"," current_sa_counts=torch.tensor([sa_counts[(state_int, a)] for a in range(4)]).to(device) #pick out counts of current state\n"," with torch.no_grad():\n"," #1.student\n"," qvals_student=agent(state[0])\n"," qvals_student-=min(qvals_student)*torch.ones(n_actions).to(device)\n"," probas=softy(iota*qvals_student) #multiply by constant iota for suitable probabilities from Q-values\n"," #2.optimism in face of uncertainty\n"," probas+=alpha/(current_sa_counts+1)\n"," #argmax\n"," action=torch.argmax(probas).item()\n"," return probas, action"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":339,"status":"ok","timestamp":1702900611752,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"sDx0TccfjHQW"},"outputs":[],"source":["#@title 9. FUNCTIONS - Optimizer for DQN\n","'''\n","Optimization function for the DQN\n","'''\n","\n","def optimize_dqn(network: DQN,optimizer: torch.optim, memory: torch.tensor, batches: 'tuple[torch.tensor, torch.tensor]', goal_found: bool, loss_norm) -> float:\n"," '''\n"," Do one optimizer step of the (teacher) DQN\n"," ---\n"," INPUT\n"," network: the DQN we are optimizing\n"," optimizer: the opimizer for the DQN\n"," memory: tensor containing the indices of all transitions that are used for this training step (there can be multiple occurences of single indices)\n"," -> those indices comprise the short- and long-term memory\n"," batches: contains a concatenation of a) the batch of all states and b) the batch of nextstates (for all possible transitions)\n"," further it contains the batch of rewards (for all possible transitions)\n"," goal_found: Indicates if the goal has been found yet\n"," loss_norm: norm for the loss function (i.e. nn.MSELoss)\n"," ---\n"," OUTPUT:\n"," loss: loss in this optimizer step\n"," '''\n","\n"," states_nextstates, rewards = batches\n","\n"," #Compute all the Q-values for all states, all nextstates and all 4 actions each\n"," qvalues=network(states_nextstates)\n"," state_action_values = torch.flatten(qvalues[:int(len(rewards)/4)])\n"," # Compute max_a' Q(s_{t+1},a') for all next states\n"," next_state_values = qvalues[int(len(rewards)/4):].max(dim=1,keepdim=True).values\n","\n"," # Compute the expected Q values\n"," expected_state_action_values = torch.flatten(gamma_bellman*next_state_values + rewards) #Bellman\n","\n","\n"," #get the occurences of each transition, square root is necessary for MSELoss\n"," trans_factors=torch.sqrt(torch.bincount(memory, minlength=len(next_state_values)))\n"," #train all 4 goal state values (for the 4 actions) only if the goal has been found\n"," if goal_found:\n"," expected_state_action_values[-4:]=2.\n"," trans_factors[-4:]=mt.sqrt(2) #put every goal action twice (long- and short-term memory)\n"," # Compute loss using the Bellman equation and MSE norm (if this is to change, have to adjust the square root above)\n"," loss = 1/mt.sqrt(len(memory)+8)*loss_norm(state_action_values*trans_factors, expected_state_action_values*trans_factors) #+8 because of 4 goal actions two times\n"," if not goal_found:\n"," expected_state_action_values[-4:]=0.\n"," state_action_values[-4:]=0.\n"," # Compute loss using the Bellman equation and MSE norm (if this is to change, have to adjust the square root above)\n"," loss = 1/mt.sqrt(len(memory))*loss_norm(state_action_values*trans_factors, expected_state_action_values*trans_factors)\n","\n","\n"," # Optimize the model\n"," optimizer.zero_grad()\n"," loss.backward()\n","\n"," '''\n"," #restricting parameter gradients to lie in the interval [-1,1]\n"," for param in network.parameters():\n"," param.grad.data.clamp_(-1, 1)\n"," '''\n","\n"," optimizer.step()\n"," return loss.item()\n","\n","\n","\n","\n","def transition_memories(init_state:int, goal_state:int, wall_states:'list[int]')->'tuple[ dict[tuple[int,int],int], tuple[torch.tensor, torch.tensor] ]':\n"," '''\n"," create practical memory for transitions that allows for savings in Q-network training\n"," when learning the task\n"," ---\n"," INPUT\n"," init_state, goal_state, wall_states: specifications of the task\n"," ---\n"," OUTPUT\n"," transition_index_dict: dictionary taking as input a state-action combination and giving out the index belonging to this particular transition\n"," batches_list: list containing two tensor, firstly the concatenated states and next_states and secondly the concatenated rewards.\n"," this complicated format makes computations faster.\n"," '''\n"," #initialize environment\n"," env = SquareGridworld(init_state,goal_state,wall_states, lava)\n"," outcomes=env.get_outcomes()\n","\n"," transition_memory = ReplayMemory(capacity=n_actions*grid_dim**2)\n"," transition_index_dict={} #each transition gets an index\n"," #create the transition batches according to the world setup\n"," i=0\n"," for s_int in range(grid_dim**2):\n"," if s_int in wall_states or s_int==goal_state:\n"," continue\n"," s=state_int_to_tuple(s_int)\n"," for a in range(n_actions):\n"," transition_index_dict[(s_int,a)]=i\n"," i+=1\n"," ns_int, r = outcomes[s_int,a]\n"," ns=state_int_to_tuple(ns_int)\n"," transition_memory.push(s, torch.tensor([[a]], device=device), ns, torch.tensor([[r]], device=device))\n"," #add the goal state transitions at the end\n"," gs=state_int_to_tuple(goal_state)\n"," for a in range(n_actions):\n"," transition_memory.push(gs, torch.tensor([[a]], device=device), gs, torch.tensor([[2.]], device=device))\n"," transition_index_dict[(goal_state,a)]=i\n"," i+=1\n","\n"," transitions = transition_memory.memory\n"," #create a big transition\n"," batch = Transition(*zip(*transitions))\n"," #and now the batches to train on for the network\n"," nextstates=torch.cat(batch.next_state)\n"," rewards=torch.cat(batch.reward)\n","\n","\n"," #second memory for only the states (don't need to run all of them through the network four times, but only one time)\n"," transition_memory = ReplayMemory(capacity=n_actions*grid_dim**2)\n"," for s_int in range(grid_dim**2):\n"," if s_int in wall_states or s_int==goal_state:\n"," continue\n"," s=state_int_to_tuple(s_int)\n"," transition_memory.push(s, torch.tensor([[0]], device=device), s, torch.tensor([[0.]], device=device))\n"," #add the goal state transitions at the end\n"," gs=state_int_to_tuple(goal_state)\n"," transition_memory.push(gs, torch.tensor([[0]], device=device), gs, torch.tensor([[0.]], device=device))\n","\n"," transitions = transition_memory.memory\n"," #create a big transition\n"," batch = Transition(*zip(*transitions))\n"," states=torch.cat(batch.state)\n"," batches_list=[torch.cat([states, nextstates]), rewards]\n","\n"," return transition_index_dict, batches_list"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":3,"status":"ok","timestamp":1702900612013,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"K5BxOxPvaoMY"},"outputs":[],"source":["#@title 10. FUNCTIONS - Plots of Q-matrices and probability matrices (from Q matrices)\n","\n","def q_probabilities_plotter(Q: torch.tensor, init_state:int, goal_state:int, wall_states:'list[int]', steps:int, title: str):\n"," '''\n"," Plot the probabilites of state occupancy after a certain number of steps in a grid plot\n"," The initial state is marked with green text color, the final state with gold text color\n"," ---\n"," INPUT\n"," Q: the q-matrix to be used for probability calculation\n"," init_state, goal_state, wall_states: properties of the gridworld\n"," steps: number of steps for the agent after which the probabilites are calculated\n"," title: title for the plot\n"," ---\n"," OUTPUT\n"," Nothing except the plot\n"," '''\n","\n"," cmap=plt.get_cmap(\"binary\")\n"," softy=nn.Softmax(dim=0)\n"," #first build the network architecture\n"," env = SquareGridworld(init_state,goal_state,wall_states, lava)\n"," outcomes = env.get_outcomes()\n"," #dictionary to retrieve next state and reward given current (s,a)\n"," next_states_dict={s:[outcomes[s,a][0] for a in range(n_actions)] for s in range(grid_dim**2)}\n"," matrix_big=torch.zeros(size=(grid_dim**2,grid_dim**2), device=device) #create a big transition matrix\n"," #get action probabilities as softmax of the Q-values\n"," action_probas=softy(Q)\n"," for s in range(grid_dim**2):\n"," if s!=goal_state:\n"," for i,ns in enumerate(next_states_dict[s]):\n"," matrix_big[ns,s]+=action_probas[i,mt.floor(s/grid_dim),s%grid_dim] #need to have the \"+=\" here because there can be multiple identical transitions, e.g. for corner states\n"," else:\n"," matrix_big[s,s]=1\n"," #initialize probabilities\n"," probas=torch.zeros(grid_dim**2, device=device)\n"," probas[init_state]=1\n"," #Apply the transition matrix to the initial probability distribution for each allowed step to get the final probability distribution\n"," for _ in range(steps):\n"," probas=matrix_big@probas\n","\n"," xydict={s:(grid_dim-1-mt.floor(s/grid_dim), s%grid_dim) for s in range(grid_dim**2)}\n"," #Create 2D lists for state probabilites and state colors\n"," state_probas=np.zeros(shape=(grid_dim,grid_dim))\n"," state_colors=np.zeros(shape=(grid_dim,grid_dim))\n"," state_probas=[list(arr) for arr in state_colors]\n"," state_colors=[list(arr) for arr in state_colors]\n"," #Insert probabilites and colors at the correct positions (so that state 0 is bottom left and state grid_dim**2-1 is top right)\n"," for s in range(grid_dim**2):\n"," x,y=xydict[s]\n"," if s in wall_states:\n"," state_probas[x][y]=f\"0%\"\n"," state_colors[x][y]=cmap(1-probas[s].item())\n"," else:\n"," state_probas[x][y]=f\"{round(100*probas[s].item(),2)}%\"\n"," state_colors[x][y]=cmap(0.75*(1-probas[s].item()))\n","\n","\n"," #Now plot as 2D plot with probabilities and state numbers\n"," fig, ax = plt.subplots()\n"," # hide axes\n"," fig.patch.set_visible(False)\n"," ax.axis('off')\n"," ax.axis('tight')\n"," table=ax.table(cellText=state_probas, cellColours=state_colors, loc='center', colWidths=[0.2]*grid_dim, cellLoc='center')\n"," #some tweaking of font and text color (initial state green, goal state gold)\n"," table._cells[xydict[init_state]]._text.set_color('darkgreen')\n"," table._cells[xydict[goal_state]]._text.set_color('goldenrod')\n"," table.scale(0.72,4)\n"," for (row, col), cell in table.get_celld().items():\n"," cell.set_text_props(fontproperties=FontProperties(size=35))\n"," fig.tight_layout()\n"," plt.title(title)\n"," plt.show()\n","\n","\n","\n","def q_arrows_plotter(Q:torch.tensor, init_state: int, goal_state:int, wall_states:'list[int]', bad_trust_states: 'list[int]',\n"," good_trust_states: 'list[int]', title: str, save_loc: str, save: bool):\n"," '''\n"," Plot the probabilites of state transitions as arrows in a grid plot\n"," The initial state is marked with green color, the final state with gold color\n"," ---\n"," INPUT\n"," Q: the q-matrix to be used for probability calculation\n"," init_state, goal_state, wall_states: properties of the gridworld\n"," bad_trust_states/good_trust_states: states in which the trust was bad/good\n"," title: title for the plot\n"," save_loc: saving location for the plot.\n"," save: True if we should save the figure, False otherwise\n"," ---\n"," OUTPUT\n"," Nothing except the plot\n"," '''\n"," plot_dim=grid_dim+2 #outside walls on each side mean the dimension is increased by 2\n"," #change wall states to incorporate outside states\n"," outer_wall_states=[i for i in range(plot_dim**2) if iplot_dim*(plot_dim-1)]\n"," outer_bt_states, outer_gt_states= [],[]\n"," for w in wall_states:\n"," w1,w2=w%grid_dim, mt.floor(w/grid_dim)\n"," outer_wall_states+=[plot_dim*w2+w1+plot_dim+1]\n"," wall_states=outer_wall_states\n"," i1,i2=init_state%grid_dim, mt.floor(init_state/grid_dim)\n"," init_state=plot_dim*i2+i1+plot_dim+1\n"," g1,g2=goal_state%grid_dim, mt.floor(goal_state/grid_dim)\n"," goal_state=plot_dim*g2+g1+plot_dim+1\n"," for b in bad_trust_states:\n"," b1,b2=b%grid_dim, mt.floor(b/grid_dim)\n"," outer_bt_states+=[plot_dim*b2+b1+plot_dim+1]\n"," bad_trust_states=outer_bt_states\n"," for g in good_trust_states:\n"," g1,g2=g%grid_dim, mt.floor(g/grid_dim)\n"," outer_gt_states+=[plot_dim*g2+g1+plot_dim+1]\n"," good_trust_states=outer_gt_states\n","\n","\n"," cmap=plt.get_cmap(\"binary\")\n"," softy=nn.Softmax(dim=0)\n"," #get action probabilities as softmax of the Q-values\n"," action_probas=softy(Q)\n","\n"," xydict={s:(plot_dim-1-mt.floor(s/plot_dim), s%plot_dim) for s in range(plot_dim**2)} #+1 to include outside walls\n"," #Create 2D lists for state probabilites and state colors\n"," state_colors=np.zeros(shape=(plot_dim,plot_dim))\n"," state_colors=[list(arr) for arr in state_colors]\n"," for s in range(plot_dim**2):\n"," x,y=xydict[s]\n"," if s in wall_states:\n"," state_colors[x][y]=(122/255, 176/255, 207/255,1)\n"," elif s in bad_trust_states:\n"," state_colors[x][y]=(255/255, 170/255, 170/255,1)\n"," elif s in good_trust_states:\n"," state_colors[x][y]=(144/255, 238/255, 144/255,1)\n"," else:\n"," state_colors[x][y]=cmap(0.)\n","\n"," '''\n"," elif s==init_state:\n"," state_colors[x][y]=(239/255, 138/255, 98/255,1)\n"," elif s==goal_state:\n"," state_colors[x][y]=(95/255, 182/255, 119/255,1)\n"," '''\n","\n"," #Now plot as 2D plot with probabilities and state numbers\n"," fig, ax = plt.subplots(figsize=(8,8))\n"," # hide axes\n"," fig.patch.set_visible(True)\n"," ax.axis('off')\n"," ax.axis('tight')\n"," table=ax.table(cellColours=state_colors, loc='center', colWidths=[0.2]*plot_dim, cellLoc='center')\n"," table.scale(1,7)\n"," cell_size_x, cell_size_y=[0.0221,0.0212] #have to somehow find out the cell size to plot proper arrows\n"," for s in range(plot_dim**2):\n"," if s in wall_states or s==goal_state:\n"," continue\n"," x,y=cell_size_x*(s%plot_dim-plot_dim/2+0.5), cell_size_y*(mt.floor(s/plot_dim)-plot_dim/2+0.5) #coordinates on the plot of our state\n"," #transform back the state index\n"," s_original=s-plot_dim-1\n"," if s>=2*plot_dim:\n"," s_original-=2*mt.floor(s/plot_dim-1)\n"," indx,indy=mt.floor(s_original/grid_dim),s_original%grid_dim\n"," probas=np.round(action_probas[:,indx,indy].detach().cpu().numpy(),2)\n"," maxdir=np.argmax(probas)\n"," #all four possible directions\n"," wbase=1.8e-3\n"," hb=2e-3\n"," hb2=3e-3\n"," if maxdir==0:\n"," plt.arrow(x+cell_size_x/4,y+cell_size_y/8,cell_size_x/2,0, width=wbase*mt.sqrt(probas[0]),length_includes_head=True, head_width=hb+hb2*mt.sqrt(probas[0]), head_length=hb+hb2*mt.sqrt(probas[0]), color=\"firebrick\") #right\n"," #plt.text(x+cell_size_x/4,y+cell_size_y/6, f\"{round(100*probas[0])}\", horizontalalignment=\"left\", verticalalignment=\"bottom\")\n"," elif maxdir==2:\n"," plt.arrow(x-cell_size_x/4,y-cell_size_y/8,-cell_size_x/2,0, width=wbase*mt.sqrt(probas[2]),length_includes_head=True, head_width=hb+hb2*mt.sqrt(probas[2]), head_length=hb+hb2*mt.sqrt(probas[2]),color=\"firebrick\") #left\n"," #plt.text(x-cell_size_x/4,y-cell_size_y/12, f\"{round(100*probas[2])}\", horizontalalignment=\"right\", verticalalignment=\"bottom\")\n"," elif maxdir==1:\n"," plt.arrow(x+cell_size_x/8,y+cell_size_y/4,0,cell_size_y/2, width=wbase*mt.sqrt(probas[1]),length_includes_head=True, head_width=hb+hb2*mt.sqrt(probas[1]), head_length=hb+hb2*mt.sqrt(probas[1]),color=\"firebrick\") #up\n"," #plt.text(x+cell_size_x/6,y+3*cell_size_y/4, f\"{round(100*probas[1])}\", horizontalalignment=\"left\", verticalalignment=\"top\")\n"," else:\n"," plt.arrow(x-cell_size_x/8,y-cell_size_y/4,0,-cell_size_y/2, width=wbase*mt.sqrt(probas[3]),length_includes_head=True, head_width=hb+hb2*mt.sqrt(probas[3]), head_length=hb+hb2*mt.sqrt(probas[3]),color=\"firebrick\") #down\n"," #plt.text(x-cell_size_x/6,y-3*cell_size_y/4, f\"{round(100*probas[3])}\", horizontalalignment=\"right\", verticalalignment=\"bottom\")\n"," plt.title(title)\n"," if save:\n"," plt.savefig(save_loc, format=\"svg\")\n"," plt.figure()\n"," if not save:\n"," plt.show()\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":420,"status":"ok","timestamp":1702900613375,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"pEFfZ0oXQQWR"},"outputs":[],"source":["#@title 11. FUNCTIONS - Graph operations\n","\n","def graph_from_walls(wall_states: 'list[int]')->nx.Graph:\n"," '''\n"," Creates a graph representing the gridworld (goal state or initial state are not highlighted)\n"," The edges have weights corresponding to step rewards or wall rewards\n"," ---\n"," INPUT\n"," wall_states: wall state positions\n"," ---\n"," OUTPUT\n"," G: graph representing the gridworld\n"," '''\n"," G=nx.Graph()\n"," #add nodes\n"," for s in range(grid_dim**2):\n"," G.add_node(s,pos=(s%grid_dim,mt.floor(s/grid_dim))) #so that the coordinate (0,0) is in the middle of the gridworld\n"," #add \"regular\" edges (no wall states)\n"," for s1 in range(grid_dim**2):\n"," for s2 in [s1+1, s1+grid_dim]: #s1+1: horizontal edges, s1+grid_dim: vertical edges\n"," if s2 in G.nodes:\n"," if (s2==s1+1 and s1%grid_dim==grid_dim-1) or (s2==s1+grid_dim and s1>=grid_dim*(grid_dim-1)): #exceptions for right and top walls - here make no connection\n"," continue\n"," else:\n"," G.add_edge(s1,s2, weight=-step_reward) #we make the weights here positive so that later shortest-path-algorithms can be applied easily\n"," if lava:\n"," #for wall states, change the edge weights leading to them and away from them\n"," for s in wall_states:\n"," for (s_i,s_j) in [[s-1,s],[s,s+1],[s-grid_dim, s],[s,s+grid_dim]]: #change all edges extending from the wall state\n"," if (s_i,s_j) in G.edges:\n"," G.edges[s_i,s_j]['weight']-=wall_reward #add the negative wall reward to these edges\n"," else:\n"," for s in wall_states:\n"," G.remove_node(s) #remove inaccessible wall states\n"," return G\n","\n","\n","def max_dist_pair(G:nx.Graph)->'tuple[int,int]':\n"," '''\n"," Compute pair of states with maximum distance between them in a given gridworld\n"," ---\n"," INPUT\n"," G: the graph representing the gridworld as created by graph_from_walls\n"," ---\n"," OUTPUT\n"," (s1,s2): The pair of states that are furthest apart form each other in the gridworld\n"," '''\n"," #find pair of nodes with maximum distance between them using dijkstra\n"," all_pairs=dict(nx.all_pairs_dijkstra(G))\n"," all_pairs_dist=[(n,all_pairs[n][0]) for n in all_pairs]\n"," max_dist_per_node=[(n,list(dijkdict.keys())[np.argmax(np.array(list(dijkdict.values())))],max(list(dijkdict.values()))) for n,dijkdict in all_pairs_dist]\n"," best_index=np.argmax([k for i,j,k in max_dist_per_node])\n"," s1, s2 = max_dist_per_node[best_index][0], max_dist_per_node[best_index][1]\n"," return s1, s2\n","\n","\n","def dead_end_goals(wall_states:'list[int]')->'list[int]':\n"," '''\n"," given a maze where lanes to walk are only a single step wide, this function calculates where possible goals can be (in dead ends)\n"," ---\n"," INPUT:\n"," wall_states: list of wall state positions\n"," ---\n"," OUTPUT:\n"," dead_ends: list of dead end states\n"," '''\n"," dead_ends=[]\n"," for s in range(grid_dim**2):\n"," #calculate neighbour \"values\" -> if -1 it means that there is a wall\n"," left_nb = -1 if s%grid_dim == 0 else s-1\n"," up_nb = -1 if s>=grid_dim*(grid_dim-1) else s+grid_dim\n"," right_nb = -1 if s%grid_dim == grid_dim-1 else s+1\n"," down_nb = -1 if sint:\n"," '''\n"," Compute shortest path length between two given states\n"," ---\n"," INPUT\n"," G: the graph representing the gridworld as created by graph_from_walls\n"," s1, s2: two integers representing the two states\n"," ---\n"," OUTPUT\n"," dist - The distance between the two nodes in steps\n"," '''\n"," #find distance between the nodes (use weight=None to count every step as distance 1)\n"," dist=nx.dijkstra_path_length(G,s1,s2, weight=None)\n"," return dist\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":2,"status":"ok","timestamp":1702900613839,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"9UYRMv2yjpIX"},"outputs":[],"source":["#@title 12. FUNCTIONS - Q matrix generation from gridworlds\n","\n","\n","def q_matrix_generator(label_dict: 'dict[int,tuple[int, int, int]]', wall_state_dict: 'dict[int,list[int]]', perfect_qdict: 'dict[int, torch.tensor]', matrix_accuracy: float, max_attempts: int,\n"," batch_size: Optional[int]=None, num_eps: int=200, loss_norm=nn.MSELoss(), max_steps: int=50, memory_size_short=L, alpha: float=16, learning_rate=lr_teacher)->'dict[int, torch.tensor]':\n"," '''\n"," For each maze task stored in the label dict, a (near-) optimal Q-matrix is created by the DQN and stored in the dictionary\n"," ---\n"," INPUT\n"," label_dict -dictionary with indices representing tasks as keys and corresponding tuple (wall_index, initial_state, goal_state)\n"," as value (wall index is key for wall_state_dict)\n"," wall_state_dict: dictionary with \"wall indices\" as keys and lists representing wall states of a particular gridworld as values\n"," perfect_qdict: the perfect q-matrices created by q_matrix_generator_deterministic (keyed by task integer), in the calculations we make sure that the teacher's q-matrix\n"," is not too far off the perfect q-matrix\n"," batch_size: batching of transitions in network learning (\"None\" means there is no batching and we do full gradient descent with all transitions)\n"," num_eps: Number of episodes the (teacher) agent gets\n"," loss_norm: The norm that is used for Q-learning training on the transitions in the memory\n"," max_steps: Maximum steps the agent gets per episode\n"," memory_size_short: Size of the short-term memory\n"," alpha: constant in the \"optimism in the face of uncertainty\"-policy\n"," learning_rate: learning rate for the optimizer of the DQN\n"," matrix_accuracy: Guideline for how close (in terms of 1-norm) the teacher's q-matrix has to be to the perfect q-matrix to be accepted\n"," ---\n"," OUTPUT\n"," qdict: dictionary with gridworld task indices as keys and corresponding q-matrices as values\n"," '''\n","\n"," qdict: 'dict[int, torch.tensor]'={}\n"," for task, [wall_index, init_state, goal_state] in label_dict.items():\n"," wall_states=wall_state_dict[wall_index]\n"," print(f\"task is {task}\")\n"," G=graph_from_walls(wall_states)\n"," #when the graph is not connected, then the wall states cut off some regular states -> we don't analyze this gridworld\n"," if not nx.is_connected(G):\n"," continue\n"," '''\n"," initialization\n"," '''\n","\n"," #initialize environment\n"," env = SquareGridworld(init_state,goal_state,wall_states, lava)\n"," outcomes=env.get_outcomes()\n"," #initialize softmax and student probabilities\n"," softy=nn.Softmax(dim=0)\n","\n"," #dictionary for faster switching between int and tuple representations of states\n"," state_int_to_tuple_dict={s:state_int_to_tuple(s) for s in range(grid_dim**2)}\n"," state_int_to_tuple_dict[None]=None\n","\n"," transition_index_dict, batches_list=transition_memories(init_state, goal_state, wall_states)\n","\n"," '''\n"," execution\n"," '''\n"," #We deem a q-matrix \"good\" if it is not too far off the \"perfect\" Q-matrix for the task (which was generated deterministically)\n"," good_qmatrix=False\n"," counter=0\n"," while counter if yes then include all four goal transitions into learning\n"," if not goalfound:\n"," goalfound=True\n"," ep_steps[ep]=t\n"," break\n","\n"," else:\n"," #add experience to the memory\n"," memory_short.append(transition_index_dict[state_int,action.item()])\n"," if sa_counts[state_int,action.item()]==0:\n"," memory_long.append(transition_index_dict[state_int,action.item()])\n","\n"," #episode is also finished if we took the maximum number of steps\n"," if t==max_steps-1:\n"," ep_steps[ep]+=max_steps\n"," # Move to the next state\n"," sa_counts[(state_int,action.item())]+=1\n"," state_int = next_state_int\n"," # Perform one step of the optimization (on the teacher network)\n"," memory=torch.tensor(list(memory_short)+list(memory_long)).to(device)\n"," '''\n"," lp2 = LineProfiler()\n"," lp2_wrapper = lp2(optimize_dqn)\n"," lp2_wrapper(student, optimizer, memory, batches_list, goalfound, loss_norm)\n"," lp2.print_stats()\n"," '''\n"," current_loss=optimize_dqn(teacher, optimizer, memory, batches_list, goalfound, loss_norm)\n","\n"," #create student Q-matrix and print out\n"," q_matrix=q_matrix_from_network(teacher, message ,wall_states)\n"," difference_to_perfect=torch.linalg.vector_norm(q_matrix-perfect_qdict[task], ord=1)\n"," print(f\"difference to perfect Q-matrix: {difference_to_perfect}\")\n"," if difference_to_perfect < matrix_accuracy:\n"," good_qmatrix=True #if False, we repeat the calculation, the final Q-matrix was not \"good enough\"\n"," qdict[task]=q_matrix\n"," q_array=np.round(np.flip(q_matrix.detach().numpy(),1).copy(),3) #flip dimension so it has the correct maze form when printing\n"," print(f\"final q-matrix of student is {q_array}\")\n"," elif counter==max_attempts-1:\n"," qdict[task]=perfect_qdict[task]\n"," counter+=1\n"," else:\n"," counter+=1\n"," #plt.plot(range(num_eps),ep_steps)\n"," plt.show()\n"," return qdict\n","\n","\n","\n","def q_matrix_from_network(network: nn.Module, message: torch.tensor, wall_states:'list[int]')->torch.tensor:\n"," '''\n"," generate the entire q-matrix of a task from the (trained) network\n"," ---\n"," INPUT\n"," network - a teacher network that gets a state as input and outputs the four corresponding q-values\n"," message - the message the network may have received\n"," wall_states - list of the wall states of the gridworld under consideration\n"," ---\n"," OUTPUT\n"," q_matrix - the q-matrix\n"," '''\n","\n"," q_matrix=torch.zeros(size=(4,grid_dim,grid_dim))\n"," for s in range(grid_dim**2):\n"," indx,indy=mt.floor(s/grid_dim),s%grid_dim\n"," if (not (s in wall_states)) or lava: #lava states are accessible, so we have trained teacher q-values for them\n"," q_matrix[:,indx,indy]=network(torch.cat((state_int_to_tuple(s)[0],message),0))\n"," else:\n"," q_matrix[:,indx,indy]=0*network(state_int_to_tuple(s)) #add zeros as placeholder values for the inaccessible wall states\n","\n"," return q_matrix\n","\n","\n","\n","def q_matrix_generator_deterministic(label_dict: 'dict[int,tuple[int, int, int]]', wall_state_dict: 'dict[int,list[int]]')->'dict[int, torch.tensor]':\n"," '''\n"," function to create (perfect) Q-matrices for maze tasks in a deterministic way - they can then be compared to the\n"," calculated Q-matrices of the teacher DQN (or student)\n"," ---\n"," INPUT\n"," label_dict: dictionary with indices representing tasks as keys and corresponding tuple (wall_index, initial_state, goal_state)\n"," as value (wall index is key for wall_state_dict)\n"," wall_state_dict: dictionary with \"wall indices\" as keys and lists representing wall states of a particular gridworld as values\n"," ---\n"," OUTPUT\n"," qdict - dictionary with indices of gridworld tasks as keys and corresponding \"perfect\" deterministic q-matrices as values\n","\n"," '''\n"," qdict={} #final outputs\n","\n"," for task, [wall_index, init_state, goal_state] in label_dict.items():\n"," wall_states=wall_state_dict[wall_index]\n"," G=graph_from_walls(wall_states)\n"," if not nx.is_connected(G):\n"," continue\n","\n"," #initialize environment\n"," env = SquareGridworld(init_state,goal_state,wall_states, lava)\n"," outcomes=env.get_outcomes()\n","\n"," #create q-matrix from the network results\n"," q_matrix=torch.zeros(size=(4,grid_dim,grid_dim))\n"," v_value_dict={} #value function for all the states\n"," v_value_dict[goal_state]=goal_reward\n","\n"," #first find the value function v(s)\n"," for s in range(grid_dim**2):\n"," if (lava or not (s in wall_states)) and not s==goal_state:\n"," reverse_goal_path=nx.dijkstra_path(G, s, goal_state, weight='weight')\n"," reverse_goal_path.reverse() #reverse the path to move backwards step by step from the goal\n"," v_value=goal_reward\n"," for k in range(len(reverse_goal_path)-1): #iterate through the edges\n"," edge_reward=-G.edges[(reverse_goal_path[k], reverse_goal_path[k+1])]['weight']\n"," v_value=gamma_bellman*v_value+edge_reward #discount every step and add the (negative) reward\n"," v_value_dict[s]=v_value\n","\n"," #then generate q(s,a) from v(s)\n"," for s in range(grid_dim**2):\n"," indx,indy=mt.floor(s/grid_dim),s%grid_dim\n"," if s in wall_states:\n"," q_matrix[:,indx,indy]=torch.tensor([0,0,0,0], device=device)\n"," elif s==goal_state:\n"," q_matrix[:,indx,indy]=torch.tensor([goal_reward, goal_reward, goal_reward, goal_reward], device=device)\n"," else:\n"," for a in range(4):\n"," next_state, reward=outcomes[s,a]\n"," q_matrix[a,indx,indy]=gamma_bellman*v_value_dict[next_state]+reward\n","\n"," qdict[task]=q_matrix\n"," print(f\"Completed calculation of deterministic Q-matrices\")\n"," return qdict\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":2,"status":"ok","timestamp":1702900614081,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"SAeTfs7u5YWY"},"outputs":[],"source":["#@title 13. FUNCTIONS - Autoencoder definition + corresponding matrix data\n","'''\n","Data Class definition (matrices) for the autoencoder\n","'''\n","\n","#This is a class for our input data, which are the q-matrices (needs to be structured like this..)\n","class MatrixDataset(Dataset):\n"," #The __init__ function is run once when instantiating the Dataset object.\n"," def __init__(self, matrix_dict:'dict[int, torch.tensor]', label_dict: 'dict[int,tuple[int, int, int]]'):\n"," self.labels=label_dict #dictionary that stores matrix labels (grid world number/init state/goal state) keyed by index\n"," self.matrices = matrix_dict #dictionary that stores matrices keyed by index\n","\n"," #The __len__ function returns the number of samples in our dataset\n"," def __len__(self)->int:\n"," return len(self.labels)\n","\n"," #The __getitem__ function loads and returns a sample from the dataset at the given index idx.\n"," def __getitem__(self, idx: int)->'tuple[torch.tensor, tuple[int,int,int]]':\n"," matrix = self.matrices[idx]\n"," labels = self.labels[idx]\n"," return matrix, labels\n","\n","\n","'''\n","convolutional autoencoder class to encode q-matrices into messages\n","'''\n","\n","\n","def conv2d_output_dims(x: 'tuple[int,int,int]', layer: nn.Conv2d)->'tuple[int,int,int]':\n"," \"\"\"\n"," Unnecessarily complicated but complete way to\n"," calculate the output depth, height\n"," and width size for a Conv2D layer\n"," ---\n"," INPUT\n"," Args:\n"," x: Input size (depth, height, width)\n"," layer: The Conv2D layer\n"," ---\n"," OUTPUT:\n"," Tuple of out-depth/out-height and out-width\n"," Output shape as given in [Ref]\n"," Ref:\n"," https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html\n"," \"\"\"\n"," assert isinstance(layer, nn.Conv2d)\n"," p = layer.padding if isinstance(layer.padding, tuple) else (layer.padding,)\n"," k = layer.kernel_size if isinstance(layer.kernel_size, tuple) else (layer.kernel_size,)\n"," d = layer.dilation if isinstance(layer.dilation, tuple) else (layer.dilation,)\n"," s = layer.stride if isinstance(layer.stride, tuple) else (layer.stride,)\n"," in_depth, in_height, in_width = x\n"," out_depth = layer.out_channels\n"," out_height = 1 + (in_height + 2 * p[0] - (k[0] - 1) * d[0] - 1) // s[0]\n"," out_width = 1 + (in_width + 2 * p[-1] - (k[-1] - 1) * d[-1] - 1) // s[-1]\n"," return (out_depth, out_height, out_width)\n","\n","\n","class ConvAutoEncoder(nn.Module):\n"," '''\n"," A Convolutional AutoEncoder\n"," '''\n"," def __init__(self, x_dim: 'tuple[int,int,int]', K: int, nonlinear_ae: bool, nonlinear_std: bool, n_filters: int=10, filter_size: int=2):\n"," '''\n"," Initialize parameters of ConvAutoEncoder\n"," ---\n"," INPUT\n"," x_dim: Input dimensions (channels, height, widths)\n"," K: message length/hidden dimension\n"," nonlinear_ae: are the activations in the autoencoder nonlinear?\n"," nonlinear_std: are the activations in the student nonlinear?\n"," n_filters: Number of filters (number of output channels)\n"," filter_size: Kernel size\n"," '''\n"," super().__init__()\n"," channels, height, widths = x_dim\n","\n"," # Encoder input bias layer\n"," self.enc_bias = BiasLayer(x_dim)\n"," # First encoder conv2d layer\n"," #32 different filters -> grid_dim x grid_dim x n_actions to grid_dim+1 x grid_dim+1 x 32\n"," self.enc_conv_1 = nn.Conv2d(channels, n_filters, filter_size, padding=filter_size-1, device=device)\n"," #32 different filters -> grid_dim+1 x grid_dim+1 x 32 to grid_dim+2 x grid_dim+2 x 32\n"," # Output shape of the first encoder conv2d layer given x_dim input\n"," conv_1_shape = conv2d_output_dims(x_dim, self.enc_conv_1)\n"," # Second encoder conv2d layer\n"," self.enc_conv_2 = nn.Conv2d(n_filters, n_filters, filter_size, padding=filter_size-1, device=device) #and here once again 32 different filters?!\n"," # Output shape of the second encoder conv2d layer given conv_1_shape input\n"," conv_2_shape = conv2d_output_dims(conv_1_shape, self.enc_conv_2)\n"," # The bottleneck is a dense layer, therefore we need a flattenning layer\n"," self.enc_flatten = nn.Flatten()\n"," # Conv output shape is (depth, height, width), so the flatten size is:\n"," flat_after_conv = conv_2_shape[0] * conv_2_shape[1] * conv_2_shape[2]\n"," # Encoder Linear layer\n"," self.enc_lin = nn.Linear(flat_after_conv, K, device=device)\n","\n"," # Decoder Linear layer\n"," self.dec_lin = nn.Linear(K, flat_after_conv, device=device)\n"," # Unflatten data to (depth, height, width) shape\n"," self.dec_unflatten = nn.Unflatten(dim=-1, unflattened_size=conv_2_shape)\n"," # First \"deconvolution\" layer\n"," self.dec_deconv_1 = nn.ConvTranspose2d(n_filters, n_filters, filter_size, padding=filter_size-1, device=device)\n"," # Second \"deconvolution\" layer\n"," self.dec_deconv_2 = nn.ConvTranspose2d(n_filters, channels, filter_size, padding=filter_size-1, device=device)\n"," # Decoder output bias layer\n"," self.dec_bias = BiasLayer(x_dim)\n","\n"," #Student layers (only linear and bias layers)\n"," self.std_lin1=nn.Linear(2+K, 10, device=device) #input size is 2(gridworld coordinates)+length of message -> adjustable later\n"," self.std_lin1_bias=BiasLayer(10)\n"," self.std_lin2=nn.Linear(10,20, device=device) #first intermediate layer\n"," self.std_lin2_bias=BiasLayer(20)\n"," self.std_lin3=nn.Linear(20,20, device=device) #second intermediate layer\n"," self.std_lin3_bias=BiasLayer(20)\n"," self.std_lin4=nn.Linear(20,4, device=device) #output layer\n"," self.std_lin4_bias=BiasLayer(4)\n","\n"," #booleans marking the nonlinearities\n"," self.nonlinear_ae=nonlinear_ae\n"," self.nonlinear_std=nonlinear_std\n","\n","\n","\n","\n","\n"," def encode(self, q:torch.tensor)->torch.tensor:\n"," '''\n"," first half of autoencoder: encode q-matrix to create the message\n"," ---\n"," INPUT\n"," q: The Q-matrix\n"," ---\n"," OUTPUT\n"," m: The message, i.e. the encoded Q-matrix\n"," '''\n"," m = self.enc_bias(q)\n","\n"," #nonlinear\n"," if self.nonlinear_ae:\n"," m = F.relu(self.enc_conv_1(m))\n"," m = F.relu(self.enc_conv_2(m))\n"," #linear\n"," else:\n"," m=self.enc_conv_1(m)\n"," m=self.enc_conv_2(m)\n","\n"," m = self.enc_flatten(m)\n"," m = self.enc_lin(m)\n"," return m\n","\n","\n"," def decode_ae(self, m:torch.tensor)->torch.tensor:\n"," '''\n"," second half of autoencoder: reconstruct the original q-matrix from the message\n"," ---\n"," INPUT\n"," m: The message\n"," ---\n"," OUTPUT\n"," q: The decoded Q-matrix\n"," '''\n","\n"," #nonlinear\n"," if self.nonlinear_ae:\n"," q = F.relu(self.dec_lin(m))\n"," q = self.dec_unflatten(q)\n"," q = F.relu(self.dec_deconv_1(q))\n"," #linear\n"," else:\n"," q=self.dec_lin(m)\n"," q = self.dec_unflatten(q)\n"," q=self.dec_deconv_1(q)\n","\n"," q = self.dec_deconv_2(q)\n"," q = self.dec_bias(q)\n"," return q\n","\n","\n","\n","\n"," def student(self, m_all: torch.tensor, state_tensors: torch.tensor)->torch.tensor:\n"," '''\n"," student decoding - the student network creates its q-matrix from the message(s)\n"," Could probably improve this function by replacing the for-loop..\n"," ---\n"," INPUT\n"," m_all: A number of messages combined in one tensor\n"," ---\n"," OUTPUT\n"," Q_all: torch.tensor. The student Q-matrices corresponding to the input messages\n"," '''\n"," Q_all=torch.zeros(size=[len(m_all),n_actions,grid_dim,grid_dim]).to(device)\n"," m_all=m_all.repeat(grid_dim,1)\n"," for j in range(grid_dim):\n"," ms=torch.cat((m_all,state_tensors[j]), dim=1)\n","\n"," #nonlinear\n"," if self.nonlinear_std:\n"," x=F.relu(self.std_lin1(ms))\n"," x=self.std_lin1_bias(x)\n"," x=F.relu(self.std_lin2(x))\n"," x=self.std_lin2_bias(x)\n"," x=F.relu(self.std_lin3(x))\n","\n"," #linear\n"," else:\n"," x=self.std_lin1(ms)\n"," x=self.std_lin1_bias(x)\n"," x=self.std_lin2(x)\n"," x=self.std_lin2_bias(x)\n"," x=self.std_lin3(x)\n","\n"," x=self.std_lin3_bias(x)\n"," x=self.std_lin4(x)\n"," x=self.std_lin4_bias(x)\n"," for i in range(grid_dim):\n"," Q_all[:,:,i,j]=x[i*int(len(m_all)/grid_dim):(i+1)*int(len(m_all)/grid_dim)]\n"," return Q_all\n","\n","\n"," def forward(self, q: torch.tensor)->'tuple[torch.tensor, torch.tensor]':\n"," '''\n"," do a forward pass of the autoencoder, i.e. encoding and decoding, but without the student\n"," ---\n"," INPUT\n"," q: A number of messages combined in one tensor\n"," OUTPUT\n"," m: The student Q-matrices corresponding to the input messages\n"," q_rec: The reconstructed Q-matrix (by the second half of the autoencoder)\n"," '''\n"," m=self.encode(q)\n"," q_rec=self.decode_ae(m)\n"," return m, q_rec\n","\n"," #do a full pass of the autoencoder, i.e. encoding/decoding (WITH student)\n"," def forward_student(self, q: torch.tensor, state_tensors: torch.tensor)->'tuple[torch.tensor, torch.tensor, torch.tensor]':\n"," '''\n"," do a forward pass of the autoencoder, i.e. encoding and decoding, and also the student\n"," ---\n"," INPUT\n"," q: A number of messages combined in one tensor\n"," OUTPUT\n"," m: The student Q-matrices corresponding to the input messages\n"," q_rec: The reconstructed Q-matrix (by the second half of the autoencoder)\n"," q_std: The Q-matrix created by the student\n"," '''\n"," m=self.encode(q)\n"," q_rec=self.decode_ae(m)\n"," '''\n"," lp = LineProfiler()\n"," lp_wrapper = lp(self.student)\n"," lp_wrapper(m)\n"," lp.print_stats()\n"," '''\n"," q_std=self.student(m, state_tensors)\n"," return m, q_rec, q_std\n","\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":2,"status":"ok","timestamp":1702900615286,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"YhhCTLDI53Nv"},"outputs":[],"source":["#@title 14. FUNCTIONS - Autoencoder training\n","'''\n","function to train the autoencoder network (includes several sub-loss functions)\n","'''\n","\n","def train_autoencoder(autoencoder: nn.Module, optim: torch.optim, dataset: MatrixDataset, wall_state_dict: 'dict[int,list[int]]', gamma_sparse: float, zeta_std: float, kappa: float,\n"," epochs: int, batch_size: Optional[int], student_incl: bool, train_order: int)-> torch.tensor:\n"," '''\n"," Function to train the autoencoder network (=language proxy)\n"," ---\n"," INPUT\n"," autoencoder - Autoencoder instance\n"," optim - the optimizer for the autoencoder network\n"," dataset - the q-matrix dataset used for training\n"," wall_state_dict - dictionary with indices as keys and lists representing wall states of a particular gridworld as values\n"," gamma_sparse - balances how much focus is put on sparsity loss compared to mse loss\n"," zeta_std - balances how much focus is put on student loss (finding the goal), compared to raw autoencoder loss\n"," kappa - balances how much focus is put on regularization, compared to the probability of finding the goal\n"," epochs - Number of training epochs\n"," batch_size - Batch size\n"," student_incl - Is the student network included in autoencoder training?\n"," train_order - 1: train only the messages with fixed student / 2: train only the student with fixed messages / 0:train both at the same time\n"," ---\n"," OUTPUT\n"," overall_losses: tensor with floats documenting training losses over time (one value per epoch)\n"," '''\n","\n"," epoch_printout: int=25 #after how many do we print out the current losses\n"," big_next_states_dict={}\n"," probas_transformer_dict={}\n","\n","\n","\n"," #some modules have to be frozen in case we only train the messaging (autoencoder) or the understanding (student)\n"," if student_incl:\n"," for i,module in enumerate(autoencoder.modules()):\n"," if i>0:\n"," if i<11:\n"," module.requires_grad_(not (train_order==2))\n"," else:\n"," module.requires_grad_(not (train_order==1))\n","\n"," #initialize data\n"," loader = DataLoader(dataset,batch_size=batch_size,shuffle=True,drop_last=False) #drop_last=False means we include all data in training, the last batch is smaller each epoch\n","\n"," #one loop through all the data to save some computing in every iteration later\n"," for matrix_batch, labels_batch in loader:\n"," for t in range(len(matrix_batch)):\n"," wall_label, init_state, goal_state=labels_batch[0][t].item(), labels_batch[1][t].item(), labels_batch[2][t].item()\n"," #get dictionaries\n"," env = SquareGridworld(init_state,goal_state,wall_state_dict[wall_label], lava)\n"," outcomes=env.get_outcomes()\n"," next_states_dict={s:[outcomes[s,a][0] for a in [0,1,2,3]] for s in range(grid_dim**2)}\n"," #dictionary to retreive next state and reward given current (s,a)\n"," big_next_states_dict[(wall_label, goal_state)]=next_states_dict\n","\n"," probas_transformer=torch.zeros(size=(n_actions*grid_dim**2, grid_dim**2*grid_dim**2))\n"," for s in range(grid_dim**2):\n"," if s==goal_state:\n"," for a in range(n_actions):\n"," probas_transformer[a*grid_dim**2+grid_dim*mt.floor(s/grid_dim)+s%grid_dim,s*grid_dim**2+s]=1\n"," else:\n"," for a, ns in enumerate(next_states_dict[s]):\n"," probas_transformer[a*grid_dim**2+grid_dim*mt.floor(s/grid_dim)+s%grid_dim,s*grid_dim**2+ns]=1\n"," probas_transformer_dict[(wall_label, goal_state)]=probas_transformer\n","\n"," #now the actual training loop\n"," overall_losses, reconstruction_losses, sparsity_losses = [], [], []\n"," for epoch in trange(epochs, desc='Epoch'):\n"," loss_total, student_loss_total, reconstruction_loss_total, sparse_loss_total = 0,0,0,0\n"," for matrix_batch, labels_batch in loader:\n"," state_tensors=get_state_tensors(len(matrix_batch))\n","\n","\n","\n"," optim.zero_grad(set_to_none=True) #set to none is very important as otherwise often some small gradients remain!!\n"," matrix_batch=matrix_batch.to(device)\n"," if student_incl:\n"," '''\n"," lp = LineProfiler()\n"," lp_wrapper = lp(autoencoder.forward_student)\n"," lp_wrapper(matrix_batch, state_tensors)\n"," lp.print_stats()\n"," '''\n"," messages, matrices_dec, matrices_std=autoencoder.forward_student(matrix_batch, state_tensors) #get decoded matrices by autoencoder\n"," else:\n"," messages, matrices_dec = autoencoder.forward(matrix_batch)\n","\n"," # Loss calculation for the autoencoder - \"batch\" stands for the loss in this batch, \"total\" for the overall loss in this epoch (sum of all batches)\n"," reconstruction_loss_batch=reconstruction_loss(matrix_batch, matrices_dec)\n"," sparse_loss_batch = sparsity_loss(messages)\n"," if student_incl:\n"," '''\n"," lp = LineProfiler()\n"," lp_wrapper = lp(goal_finding_loss)\n"," lp_wrapper(matrices_std, labels_batch, wall_state_dict, kappa, big_next_states_dict, probas_transformer_dict)\n"," lp.print_stats()\n"," '''\n"," student_loss_batch=goal_finding_loss(matrices_std, labels_batch, wall_state_dict, kappa, big_next_states_dict, probas_transformer_dict)\n"," #student_loss_batch=reward_loss(matrices_std, labels_batch, wall_state_dict, kappa)\n"," #overall loss is a combination\n"," if student_incl:\n"," loss_batch=(1-gamma_sparse)*reconstruction_loss_batch+gamma_sparse*sparse_loss_batch+zeta_std*student_loss_batch\n"," #loss_batch=gamma_sparse*sparse_loss_batch+zeta_std*student_loss_batch\n"," if not student_incl:\n"," loss_batch=(1-gamma_sparse)*reconstruction_loss_batch+gamma_sparse*sparse_loss_batch\n","\n"," #add losses of the current batch to the total of this epoch\n"," loss_total+=loss_batch.detach()\n"," reconstruction_loss_total+=reconstruction_loss_batch.detach()\n"," sparse_loss_total+=sparse_loss_batch.detach()\n"," if student_incl:\n"," student_loss_total+=student_loss_batch.detach()\n","\n"," loss_batch.backward() #changing of the gradients according to the loss\n"," optim.step()\n","\n"," overall_losses.append(loss_total) #list of autoencoder losses per epoch\n"," reconstruction_losses.append(reconstruction_loss_total)\n"," sparsity_losses.append(sparse_loss_total)\n"," #printouts of intermediate losses\n"," if epoch %epoch_printout==0:\n"," print(f\"losses in epoch {epoch}:\")\n"," if student_incl:\n"," print(f\"losses are sparse:{gamma_sparse*sparse_loss_total}, reconstruction:{(1-gamma_sparse)*reconstruction_loss_total}, student:{zeta_std*student_loss_total}\")\n"," if not student_incl:\n"," print(f\"losses are sparse:{gamma_sparse*sparse_loss_total}, reconstruction:{(1-gamma_sparse)*reconstruction_loss_total}\")\n"," #transfer to cpu -Why?\n"," autoencoder.to('cpu')\n"," return torch.tensor(overall_losses).cpu(), torch.tensor(reconstruction_losses).cpu(), torch.tensor(sparsity_losses).cpu()\n","\n","\n","def reconstruction_loss(matrices: torch.tensor, matrices_dec: torch.tensor)->torch.tensor:\n"," '''\n"," Loss for reproduction accuracy of input for the autoencoder\n"," ---\n"," INPUT\n"," matrices - torch.tensor. input q-matrices\n"," matrices_dec - torch.tensor. reconstructed q-matrices (from the messages)\n"," ---\n"," OUTPUT\n"," loss - the loss\n"," '''\n"," loss = torch.norm(matrices_dec-matrices,2)\n"," return loss\n","\n","def sparsity_loss(messages: torch.tensor)->torch.tensor:\n"," '''\n"," Loss for sparsity of the messages\n"," ---\n"," INPUT\n"," messages - torch.tensor. the messages\n"," ---\n"," OUTPUT\n"," loss - the loss\n"," '''\n"," loss = torch.norm(messages,1)\n"," return loss\n","\n","\n","\n","\n","def goal_finding_loss(matrices_std_batch: torch.tensor, labels_batch: torch.tensor, wall_state_dict: 'dict[int,list[int]]', kappa: float, big_next_states_dict, probas_transformer_dict)->float:\n"," '''\n"," Loss for how accurately the student finds the goal given the message\n"," ---\n"," INPUT\n"," matrices_std - batch of q-matrices of the student\n"," labels - batch of labels from the labels dictionary corresponding to the batch of student matrices\n"," wall_state_dict - dictionary with indices as keys and lists representing wall states of a particular gridworld as values\n"," kappa - balances how much focus is put on regularization, compared to the probability of finding the goal\n"," ---\n"," OUTPUT\n"," loss - the loss\n"," '''\n"," student_loss=torch.zeros(size=(1,), dtype=float).to(device)[0]\n"," #iterate through all the tasks\n"," print_output=[] #this is what we print out each iteration\n"," softy=nn.Softmax(dim=0)\n"," for t in range(len(matrices_std_batch)):\n"," #load data\n"," Q=matrices_std_batch[t]\n"," wall_label, init_state, goal_state, opt_steps=labels_batch[0][t].item(), labels_batch[1][t].item(), labels_batch[2][t].item(), labels_batch[3][t].item()\n"," wall_states=wall_state_dict[wall_label]\n"," #first build the network architecture\n"," next_states_dict=big_next_states_dict[(wall_label, goal_state)]\n"," matrix_big=torch.zeros(size=(grid_dim**2,grid_dim**2), device=device) #create a big transition matrix\n"," #get action probabilities as softmax of the Q-values\n"," action_probas=torch.flatten(softy(Q)).unsqueeze(dim=0)\n","\n"," #print(wall_states, goal_state)\n"," matrix_big=action_probas@probas_transformer_dict[(wall_label, goal_state)]\n"," '''\n"," for i in range(64):\n"," for j in range(256):\n"," hi=probas_transformer_dict[(wall_label, goal_state)][i][j]\n"," if hi>0:\n"," print(i,mt.floor(j/16), j%16)\n","\n"," sys.exit()\n"," '''\n"," matrix_big=torch.transpose(torch.unflatten(input=matrix_big, dim=1, sizes=(grid_dim**2, grid_dim**2)).squeeze(),0,1)\n","\n"," #For each student step, Apply the transition matrix to the initial probability distribution to get the new probability distribution\n"," goal_proba=torch.linalg.matrix_power(matrix_big,opt_steps)[goal_state,init_state]\n","\n"," student_loss+=(1-kappa)*(1-goal_proba)**4 #first part - probability to find the goal (difference to 1)\n"," student_loss+=kappa/mt.sqrt(grid_dim**2*n_actions)*torch.norm(Q,2) #second part - keep overall Q-values low to avoid some local minima and enhance overall stability\n","\n"," return student_loss\n","\n","\n","\n","def reward_loss(matrices_std_batch: torch.tensor, labels_batch: torch.tensor, wall_state_dict: 'dict[int,list[int]]', kappa: float)->float:\n"," '''\n"," Loss for how accurately the student finds the goal given the message\n"," ---\n"," INPUT\n"," matrices_std - batch of q-matrices of the student\n"," labels - batch of labels from the labels dictionary corresponding to the batch of student matrices\n"," wall_state_dict - dictionary with indices as keys and lists representing wall states of a particular gridworld as values\n"," kappa - balances how much focus is put on regularization, compared to the probability of finding the goal\n"," ---\n"," OUTPUT\n"," loss - the loss\n"," '''\n"," reward_loss=torch.zeros(size=(1,), dtype=float).to(device)[0]\n"," #iterate through all the tasks\n"," print_output=[] #this is what we print out each iteration\n"," softy=nn.Softmax(dim=0)\n"," for t in range(len(matrices_std_batch)):\n"," #load data\n"," Q=matrices_std_batch[t]\n"," wall_label, init_state, goal_state, opt_steps=labels_batch[0][t].item(), labels_batch[1][t].item(), labels_batch[2][t].item(), labels_batch[3][t].item()\n"," wall_states=wall_state_dict[wall_label]\n"," #first build the network architecture\n"," env = SquareGridworld(init_state,goal_state,wall_states, lava)\n"," outcomes = env.get_outcomes()\n"," #dictionary to retreive next state and reward given current (s,a)\n"," next_states_rewards_dict={s:[outcomes[s,a] for a in [0,1,2,3]] for s in range(grid_dim**2)}\n"," matrix_big=torch.zeros(size=(grid_dim**2,grid_dim**2), device=device) #create a big transition matrix\n"," #get action probabilities as softmax of the Q-values\n"," action_probas=softy(Q)\n"," state_rewards=torch.zeros(grid_dim**2, device=device)\n"," for s in range(grid_dim**2):\n"," if s!=goal_state:\n"," for i,[ns,r] in enumerate(next_states_rewards_dict[s]):\n"," current_action_probas=action_probas[i,mt.floor(s/grid_dim),s%grid_dim]\n"," state_rewards[s]+=current_action_probas*r\n"," matrix_big[ns,s]+=current_action_probas #need to have the \"+=\" here because there can be multiple identical transitions, e.g. for corner states\n"," else:\n"," matrix_big[s,s]=1\n","\n"," #initialize state occupancy probabilities\n"," probas=torch.zeros(grid_dim**2, device=device)\n"," probas[init_state]=1\n"," reward=0\n","\n"," #For each student step, Apply the transition matrix to the initial probability distribution to get the new probability distribution\n"," for rep in range(opt_steps): #strict: use the opt_steps here as well!\n"," reward+=torch.dot(probas,state_rewards) #for each step, add the average obtained reward in this step\n"," probas=matrix_big@probas\n"," reward+=probas[goal_state]*goal_reward #finally add the goal reward\n","\n"," #put worst and best possible reward for the student and from the actual obtained reward create a \"probability\", like in the goal probability loss\n"," '''\n"," best_reward=opt_steps*step_reward+reward\n"," worst_reward=opt_steps*(wall_reward+step_reward)\n"," reward_proba=(reward-worst_reward)/(best_reward-worst_reward)\n"," reward_loss+=(1-kappa)*(1-reward_proba)**4 #same method as for goal finding probability\n"," '''\n"," reward_loss+=(1-kappa)*(goal_reward+opt_steps*step_reward-reward)**2 #first part - maximise obtained reward\n"," reward_loss+=kappa/mt.sqrt(grid_dim**2*n_actions)*torch.norm(Q,2) #second part - keep overall Q-values low to avoid some local minima and enhance overall stability\n"," print_output+=[round(goal_reward+opt_steps*step_reward-reward.item(),2)] #optimize on probability but still print the lost reward output\n"," print(print_output)\n"," return reward_loss\n","\n","\n","def plot_losses(losses: list):\n"," '''\n"," plotting function for the autoencoder training losses (per epoch)\n"," losses are plotted on a log scale\n"," ---\n"," INPUT\n"," losses: Log of overall losses during the process of the model\n"," ---\n"," OUTPUT\n"," Nothing but the plot\n"," '''\n"," plt.figure()\n"," plt.plot(losses)\n"," plt.legend(['Autoencoder training loss progression'])\n"," plt.xlabel('Training batch number')\n"," plt.ylabel('Losses')\n"," plt.show()\n","\n","\n","def task_indices_sorter(label_dict, train_worlds: 'list[int]', train_goals: 'list[int]')->'list[list[int]]':\n"," '''\n"," Sorts indices of the tasks according to whether the agent knows the world or the goal\n"," ---\n"," INPUT\n"," label_dict - dictionary with indices (of gridworld \"tasks\") as keys and corresponding tuple (wall_index, initial_state, goal_state)\n"," as value (wall index comes from the wall_state_dict)\n"," train_worlds - indices of worlds that the network should train on\n"," train_goals - indices of goal locations the network should train on (only in the worlds specified by train_worlds!)\n"," ---\n"," OUTPUT\n"," goal_world_indices - indices of tasks that the network was trained on\n"," goal_noworld_indices - indices of tasks where the network was trained on the goal location, but not on the world\n"," nogoal_world_indices - indices of tasks where the network was trained on the world, but not on the goal location\n"," nogoal_noworld_indices - indices of tasks where the network was trained on neither goal location nor world\n"," '''\n","\n"," #indices marking tasks where the agent has seen the goal location or the world before (or both or neither)\n"," goal_world_inds=[]\n"," goal_noworld_inds=[]\n"," nogoal_world_inds=[]\n"," nogoal_noworld_inds=[]\n"," #iterate over all possible tasks\n"," for task_index, [world,init,goal,mindist] in label_dict.items():\n"," if goal==student_init: #do not want to train or evaluate student performance on tasks that are always solved trivially\n"," continue\n"," if world in train_worlds:\n"," if goal in train_goals:\n"," goal_world_inds+=[task_index]\n"," else:\n"," nogoal_world_inds+=[task_index]\n"," else:\n"," if goal in train_goals:\n"," goal_noworld_inds+=[task_index]\n"," else:\n"," nogoal_noworld_inds+=[task_index]\n","\n"," return [goal_world_inds, goal_noworld_inds, nogoal_world_inds, nogoal_noworld_inds]"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":647,"status":"ok","timestamp":1702900616198,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"haAiCsJ6qG-E"},"outputs":[],"source":["#@title 15. FUNCTIONS - For message plots\n","\n","def variance_analyzer(message_list: 'list[np.array]', label_list: 'list[tuple[int,int,int]]', plot_worlds)-> 'tuple[float, float, float,float]':\n"," '''\n"," find out the variance contributions of goal location and world identity (and maybe initial action or other quantities)\n"," we compare within-group variance to between-group variance - use ANOVA (ANALYSIS OF VARIANCE)\n"," 1.between-group-variance computes variance of the mean message values of groups\n"," 2.within-group-variance computes variance of messages within groups and sums them up\n"," Take care of the case where the groups might have different standard deviations but the same mean, then between-group variance is low despite the groups\n"," being very different\n"," ---\n"," INPUT\n"," message_list - all messages we are considering\n"," label_list - the task labels corresponding to the messages we are considering\n"," plot_worlds - list of all plotted worlds (with indices)\n"," ---\n"," OUTPUT\n"," relvart_worlds - relative variance (between/between+within) when considering \"world groups\"\n"," relvart_goals - relative variance (between/between+within) when considering \"goal groups\"\n"," relvart_actions - relative variance (between/between+within) when considering \"action groups\"\n"," discarded_pctg - percentage of discarded messages for action variance (those where probability of walking into wall at beginning is above 20%)\n"," '''\n","\n"," #sort the messages into groups (messages within a group are tasks from the same gridworld/with identical goal locations)\n"," #in both cases, we exclude world 0 from the analysis -> goal 0 was already excluded in Q-matrix generation\n"," world_groups_messages=[[message_list[i] for i,label in enumerate(label_list) if label[0]==j if label[2]!=student_init] for j in plot_worlds if j!=0]\n"," goal_groups_messages=[[message_list[i] for i,label in enumerate(label_list) if label[2]==j and label[0]!=0] for j in range(grid_dim**2) if j!=student_init]\n","\n"," #number of groups\n"," K_world=len(world_groups_messages)\n"," K_goal=len(goal_groups_messages)\n"," #number of total messages\n"," N_world=sum([len(i) for i in world_groups_messages])\n"," N_goal=sum([len(i) for i in goal_groups_messages])\n","\n"," print(f\"We include {N_world} messages in the {K_world} groups in the variance analysis with world groupings\")\n"," print(f\"We include {N_goal} messages in the {K_goal} groups in the variance analysis with goal groupings\")\n","\n"," #remove empty groups\n"," world_groups_messages=[group for group in world_groups_messages if len(group)!=0]\n"," goal_groups_messages=[group for group in goal_groups_messages if len(group)!=0]\n","\n"," #calculate group means\n"," world_message_means=[[len(group),np.mean(group, axis=0)] for group in world_groups_messages]\n"," goal_message_means=[[len(group),np.mean(group, axis=0)] for group in goal_groups_messages]\n","\n"," #calculate overall mean\n"," world_mean_message=np.sum([i*j for i,j in world_message_means], axis=0)/N_world\n"," goal_mean_message=np.sum([i*j for i,j in world_message_means], axis=0)/N_goal\n","\n","\n"," #between group variation\n"," var_between_worlds=np.sum([i*(np.linalg.norm(j-world_mean_message,2)**2) for i,j in world_message_means])\n"," var_between_goals=np.sum([i*(np.linalg.norm(j-goal_mean_message,2)**2) for i,j in goal_message_means])\n","\n"," #within group variation\n"," var_within_worlds=np.sum([np.linalg.norm(j-world_message_means[i][1],2)**2 for i,group in enumerate(world_groups_messages) for j in group])\n"," var_within_goals=np.sum([np.linalg.norm(j-goal_message_means[i][1],2)**2 for i,group in enumerate(goal_groups_messages) for j in group])\n","\n"," #calculate relative variations (beta)-> ratio of between-group variation to the sum of between-group and within-group variation\n"," beta_world=round(var_between_worlds/(var_within_worlds+var_between_worlds),3)\n"," beta_goal=round(var_between_goals/(var_within_goals+var_between_goals),3)\n","\n"," #calculate F-values:\n"," #first mean squares between and within groups\n"," MSb_world, MSb_goal=var_between_worlds/(K_world-1), var_between_goals/(K_goal-1)\n"," MSw_world, MSw_goal=var_within_worlds/(N_world-K_world), var_within_goals/(N_goal-K_goal)\n"," #finally get F-value\n"," F_world, F_goal=MSb_world/MSw_world, MSb_goal/MSw_goal\n","\n"," print(f\"World groups: var. within {round(var_within_worlds,2)}, var. between {round(var_between_worlds,2)}, beta {beta_world}, F {round(F_world,2)}\")\n"," print(f\"Goal groups: var. within {round(var_within_goals,2)}, var. between {round(var_between_goals,2)}, beta {beta_goal}, F {round(F_goal,2)}\")\n"," print(\"\")\n","\n","\n","\n","\n","\n","def pca_plotting(autoencoder: nn.Module, message_list: 'list[np.array]', label_list: 'dict[int,tuple[int, int, int]]', language_code: str, world_add_on: str, cmap, markersize):\n"," '''\n"," Does a PCA on a list of messages and creates two plots\n"," One barplot for the variance explained by PC\n"," One scatterplot for the messages, where PC1 vs PC2 is plotted and the messages are colored according to goal location\n"," ---\n"," INPUT\n"," autoencoder - the language proxy network generating Q-matrices from messages\n"," message_list - list of messages to consider/plot\n"," label_list - list of tuples (wall_index, initial_state, goal_state) corresponding to the messages (wall index from wall state dict)\n"," language_add_on - Marks the language from which we are analyzing the messages -> the corresponding folder where we store the figures\n"," world_add_on - Marks in the plot title if we are considering all worlds or only one specific world\n"," cmap - 2d colormap\n"," ---\n"," OUTPUT\n"," The plot.\n"," '''\n","\n","\n"," pca=PCA(K)\n"," pca.fit(message_list)\n"," message_pca=pca.fit_transform(message_list)\n","\n","\n","\n","\n"," #FIGURE WITH BARS SHOWING VARIANCE EXPLAINED BY PC\n"," axt = plt.gca()\n"," axt.tick_params(width=5, length=10)\n"," #Barplot of variance explained for all messages\n"," var_arr=np.round(pca.explained_variance_ratio_*100,2)\n"," plt.bar(range(min(K+1,11))[1:],var_arr, color=\"gray\") #for each message dimension have one PC\n"," plt.xlabel(f\"PC index\")\n"," plt.ylabel(r\"Variance\" \"\\n\" r\"explained (%)\")\n"," plt.xticks(range(min(K+1,11))[1:])\n"," plt.yticks([0,25,50,75,100])\n"," plt.ylim(0,100)\n"," print(f\"explained variances by the first {min(K,10)} PCs are {var_arr[:min(K,10)]}\")\n","\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," if save_message_plots:\n"," if world_add_on==\"all\":\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/Variance explained all messages.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/Variance explained all messages.png\", format=\"png\")\n"," else:\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/Variance explained messages world(s) {world_add_on}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/Variance explained messages world(s) {world_add_on}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","\n","\n"," #FIGURE WITH COLOR BY GOAL LOCATION IN 2D\n"," fig,ax = plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," for i, label in enumerate(label_list):\n"," goal=label[2]\n"," #transform to goal coordinates\n"," cval=(grid_dim-1)/2 #center (0,0) is in the middle of the grid\n"," goal_coords=goal%grid_dim-cval, mt.floor(goal/grid_dim)-cval\n"," #plot datapoint\n"," ax.scatter(message_pca[i,0],message_pca[i,1],color=cmap(goal_coords[0],goal_coords[1])/255, s=markersize)\n","\n"," #make some adjustments to the axes so the legend is nicely placed and doesn't overlap with data points\n"," ax.set_xlabel(f\"first PC\")\n"," xdatamin,xdatamax=min(message_pca[:,0]), max(message_pca[:,0])\n"," ydatamin,ydatamax=min(message_pca[:,1]), max(message_pca[:,1])\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n","\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," plt.title(\"Color: goal location\", y=1.05)\n"," if save_message_plots:\n"," if world_add_on==\"all\":\n"," ax.tick_params(left=False, labelleft=False) #remove ticks and labels from y axis\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA all messages, color goal.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA all messages, color goal.png\", format=\"png\")\n"," else:\n"," ax.set_ylabel(f\"second PC\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA messages, color goal, world(s) {world_add_on}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA messages, color goal, world(s) {world_add_on}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","\n","\n"," if not language_code.__contains__(\"nostudent\"):\n"," #FIGURE WITH COLOR BY INITIAL ACTION IN 2D\n"," fig,ax = plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," softy=nn.Softmax(dim=0)\n"," for i,message in enumerate(message_list):\n"," state_tensors=get_state_tensors(1)\n"," Q=autoencoder.student(torch.tensor([message]), state_tensors)[0]\n"," probas=softy(Q[:,0%grid_dim,mt.floor(0/grid_dim)]).detach().cpu().numpy()\n"," extr=0.5+grid_dim/2-1\n"," comb2d=probas[0]*np.array([extr,extr])+probas[1]*np.array([extr,-extr])+probas[2]*np.array([-extr,-extr])+probas[3]*np.array([-extr,extr])\n"," ax.scatter(message_pca[i,0],message_pca[i,1],color=cmap(comb2d[0],comb2d[1])/255, s=markersize)\n","\n"," plt.xlabel(f\"first PC\")\n"," #plt.ylabel(f\"second PC\")\n"," plt.tick_params(left=False, labelleft=False) #remove ticks and labels from y axis\n"," xdatamin,xdatamax=min(message_pca[:,0]), max(message_pca[:,0])\n"," ydatamin,ydatamax=min(message_pca[:,1]), max(message_pca[:,1])\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," plt.title(f\"Color: first student action\", y=1.05)\n"," if save_message_plots:\n"," if world_add_on==\"all\":\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA all messages, color action.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA all messages, color action.png\", format=\"png\")\n"," else:\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA messages, color action, world(s) {world_add_on}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA messages, color action, world(s) {world_add_on}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","\n","\n","\n","def tsne_plotting(autoencoder: nn.Module, perplexity: int, max_iter: int, rdkernel: int, message_list: 'list[np.array]', label_list: 'dict[int,tuple[int, int, int]]', language_code: str,\n"," world_add_on: str, scaling: bool, cmap, markersize):\n"," '''\n"," Does a t-SNE on a list of messages and creates a plot, namely\n"," a 2D scatterplot for the transformed messages, which are colored according to goal location\n"," ---\n"," INPUT\n"," autoencoder - the language proxy network generating Q-matrices from messages\n"," perplexity - something like number of neighbours to consider for the tsne-algorithm (need integer here instead of float or some weird format error in plotting will occur)\n"," max_iter - number of iterations for tsne algorithm\n"," rdkernel - random seed\n"," message_list - list of messages to consider/plot\n"," label_list - list of tuples (wall_index, initial_state, goal_state) corresponding to the messages (wall index from wall state dict)\n"," language_add_on - Marks the language from which we are analyzing the messages -> the corresponding folder where we store the figures\n"," world_add_on - Marks in the plot title if we are considering all worlds or only one specific world\n"," scaling - if True, scaling is performed before tsne application (only recommended if all dimensions should be treated with same importance)\n"," cmap - 2d colormap\n"," ---\n"," OUTPUT\n"," The plot.\n"," '''\n","\n"," tsne = TSNE(perplexity=perplexity, max_iter=max_iter, random_state=rdkernel, init=\"pca\", learning_rate=\"auto\")\n"," if scaling:\n"," message_list=StandardScaler().fit_transform(message_list)\n"," message_tsne = tsne.fit_transform(np.array(message_list))\n","\n","\n","\n","\n"," #FIGURE WITH COLOR BY GOAL LOCATION IN 2D\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," for i, label in enumerate(label_list):\n"," goal=label[2]\n"," #transform to goal coordinates\n"," cval=(grid_dim-1)/2 #center (0,0) is in the middle of the grid\n"," goal_coords=goal%grid_dim-cval, mt.floor(goal/grid_dim)-cval\n"," #plot datapoint\n"," ax.scatter(message_tsne[i,0],message_tsne[i,1],color=cmap(goal_coords[0],goal_coords[1])/255, s=markersize)\n","\n"," xdatamin,xdatamax=min(message_tsne[:,0]), max(message_tsne[:,0])\n"," ydatamin,ydatamax=min(message_tsne[:,1]), max(message_tsne[:,1])\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," plt.title(f\"Color: goal location, \"+r\"$\\pi$=\"+f\"{perplexity}\", y=1.05)\n"," if save_message_plots:\n"," if world_add_on==\"all\":\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE all messages, color goal, perplexity {perplexity}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE all messages, color goal, perplexity {perplexity}.png\", format=\"png\")\n"," else:\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE messages world(s) {world_add_on}, color goal, perplexity {perplexity}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE messages world(s) {world_add_on}, color goal, perplexity {perplexity}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","\n"," if not language_code.__contains__(\"nostudent\"):\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," #FIGURE WITH COLOR BY INITIAL ACTION IN 2D\n"," softy=nn.Softmax(dim=0)\n"," for i,message in enumerate(message_list):\n"," state_tensors=get_state_tensors(1)\n"," Q=autoencoder.student(torch.tensor([message]), state_tensors)[0]\n"," probas=softy(Q[:,student_init%grid_dim,mt.floor(student_init/grid_dim)]).detach().cpu().numpy()\n"," extr=0.5+grid_dim/2-1\n"," #combine the four probabilities into a 2-dimensional vector using the four corners of a square with four different colours\n"," comb2d=probas[0]*np.array([extr,extr])+probas[1]*np.array([extr,-extr])+probas[2]*np.array([-extr,-extr])+probas[3]*np.array([-extr,extr])\n"," ax.scatter(message_tsne[i,0],message_tsne[i,1],color=cmap(comb2d[0],comb2d[1])/255, s=markersize)\n","\n"," xdatamin,xdatamax=min(message_tsne[:,0]), max(message_tsne[:,0])\n"," ydatamin,ydatamax=min(message_tsne[:,1]), max(message_tsne[:,1])\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," plt.title(f\"Color: first student action, \"+r\"$\\pi$=\"+f\"{perplexity}\", y=1.05)\n"," if save_message_plots:\n"," if world_add_on==\"all\":\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE all messages, color action, perplexity {perplexity}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE all messages, color action, perplexity {perplexity}.png\", format=\"png\")\n"," else:\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE messages world(s) {world_add_on}, color action, perplexity {perplexity}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE messages world(s) {world_add_on}, color action, perplexity {perplexity}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","\n","\n","\n","\n","\n","\n","def umap_plotting(autoencoder: nn.Module, neighbors: int, min_dist: float, rdkernel: int, message_list: 'list[np.array]', label_list: 'dict[int,tuple[int, int, int]]', language_code: str,\n"," world_add_on: str, cmap, markersize):\n"," '''\n"," Does a Umap on a list of messages and creates a plot, namely\n"," a 2D scatterplot for the transformed messages, which are colored according to goal location\n"," ---\n"," INPUT\n"," autoencoder - the language proxy network generating Q-matrices from messages\n"," neighbors - parameter of UMAP that controls tradeoff between local and global properties\n"," min_dist - minimum distance between two points in the final embedding (prevent clumping if >0)\n"," rdkernel - random seed\n"," message_list - list of messages to consider/plot\n"," label_list - list of tuples (wall_index, initial_state, goal_state) corresponding to the messages (wall index from wall state dict)\n"," language_add_on - Marks the language from which we are analyzing the messages -> the corresponding folder where we store the figures\n"," world_add_on - Marks in the plot title if we are considering all worlds or only one specific world\n"," cmap - 2d colormap\n"," ---\n"," OUTPUT\n"," The plot.\n"," '''\n","\n"," umap_red = umap.UMAP(n_neighbors=neighbors, min_dist=min_dist, random_state=rdkernel)\n"," message_umap = umap_red.fit_transform(np.array(message_list))\n","\n","\n"," #FIGURE WITH COLOR BY GOAL LOCATION IN 2D\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," for i, label in enumerate(label_list):\n"," goal=label[2]\n"," #transform to goal coordinates\n"," cval=(grid_dim-1)/2 #center (0,0) is in the middle of the grid\n"," goal_coords=goal%grid_dim-cval, mt.floor(goal/grid_dim)-cval\n"," #plot datapoint\n"," ax.scatter(message_umap[i,0],message_umap[i,1],color=cmap(goal_coords[0],goal_coords[1])/255, s=markersize)\n","\n"," xdatamin,xdatamax=min(message_umap[:,0]), max(message_umap[:,0])\n"," ydatamin,ydatamax=min(message_umap[:,1]), max(message_umap[:,1])\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," plt.title(f\"Color: goal location, k={neighbors}\", y=1.05)\n"," if save_message_plots:\n"," if world_add_on==\"all\":\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP all messages, color goal, neighbors {neighbors}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP all messages, color goal, neighbors {neighbors}.png\", format=\"png\")\n"," else:\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP messages world(s) {world_add_on}, color goal, neighbors {neighbors}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP messages world(s) {world_add_on}, color goal, neighbors {neighbors}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","\n","\n"," if not language_code.__contains__(\"nostudent\"):\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," #FIGURE WITH COLOR BY INITIAL ACTION IN 2D\n"," softy=nn.Softmax(dim=0)\n"," for i,message in enumerate(message_list):\n"," state_tensors=get_state_tensors(1)\n"," Q=autoencoder.student(torch.tensor([message]), state_tensors)[0]\n"," probas=softy(Q[:,student_init%grid_dim,mt.floor(student_init/grid_dim)]).detach().cpu().numpy()\n"," extr=0.5+grid_dim/2-1\n"," #combine the four probabilities into a 2-dimensional vector using the four corners of a square with four different colours\n"," comb2d=probas[0]*np.array([extr,extr])+probas[1]*np.array([extr,-extr])+probas[2]*np.array([-extr,-extr])+probas[3]*np.array([-extr,extr])\n"," ax.scatter(message_umap[i,0],message_umap[i,1],color=cmap(comb2d[0],comb2d[1])/255, s=markersize)\n","\n"," xdatamin,xdatamax=min(message_umap[:,0]), max(message_umap[:,0])\n"," ydatamin,ydatamax=min(message_umap[:,1]), max(message_umap[:,1])\n"," ymean=(ydatamax+ydatamin)/2\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," plt.title(f\"Color: first student action, k={neighbors}\", y=1.05)\n"," if save_message_plots:\n"," if world_add_on==\"all\":\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP all messages, color action, neighbors {neighbors}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP all messages, color action, neighbors {neighbors}.png\", format=\"png\")\n"," else:\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP messages world(s) {world_add_on}, color action, neighbors {neighbors}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP messages world(s) {world_add_on}, color action, neighbors {neighbors}.png\", format=\"png\")\n"," plt.show()\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":2,"status":"ok","timestamp":1702900616796,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"PUXbGbbEx0NY"},"outputs":[],"source":["#@title 16. FUNCTIONS - Student solving rates\n","def student_performance_evaluator(language_code, saving_folder, language_nr, method, random_repeat: int, step_rate_bool: bool, solving_steps, rdwalker_base_rate, stepfactor, label_dict, q_matrix_dict, evaluate_nonlinear_ae, evaluate_nonlinear_std):\n"," '''\n"," Evaluates a language on the performance of the student by generating solving rates for all the possible tasks\n"," Three students are investigates:\n"," -informed student that was trained and gets the correct message\n"," -misinformed student that was trained and gets a random message from the message dictionary\n"," -untrained student that was not trained, but gets the correct message\n"," ---\n"," INPUT\n"," language_code - String. Identifier of the autoencoder network (language proxy) in the stored data\n"," saving_folder - String. Identifier of the folder where the solving rates should be saved\n"," language_nr - integer. How many languages were generated for this set of trained worlds and trained tasks (to reduce variance)\n"," method - either \"no_learning\" or \"simple_learning\"\n"," stepfactor - integer. how many times the number of steps in the optimal solution do we give the student before the episode is ended\n"," random_repeat - how many times is the random procedure for uninformed/misinformed students repeated?\n"," solving_steps - list of how many steps the random walker needed to obtain a certain solving rate - use this as baseline\n"," step_rate_bool - which method do we do, either the stepfactor method (True) or the random walker rate method (False) (for determining how many steps are allowed)\n"," evaluate_nonlinear_ae, evaluate_nonlinear_std - do the neurons in the autoencoder/the student network have nonlinear ReLU activations?\n"," ---\n"," OUTPUT\n"," None. (Store data of solving rates for all three students in a corresponding file.\n"," '''\n"," softy=nn.Softmax(dim=0)\n"," probas=torch.zeros(grid_dim**2, device=device)\n"," language_code=f\"{language_code}_language{language_nr}\"\n"," saving_folder=f\"{saving_folder}/{saving_folder}_language{language_nr}\"\n"," #iterate over languages that should be analysed\n","\n","\n"," #Load the autoencoder network\n"," autoencoder = ConvAutoEncoder(data_shape, K, evaluate_nonlinear_ae, evaluate_nonlinear_std)\n"," autoencoder.load_state_dict(torch.load(file_loc+\"autoencoder/autoencoder network parameters/\"f\"params_autoenc{language_code}.pt\", weights_only=True))\n"," autoencoder.eval()\n","\n"," #create a message dictionary, with indices corresponding to the task indices\n"," message_dict={}\n"," for task_index,q_matrix in enumerate(q_matrix_dict.values()):\n"," q_matrix=torch.unsqueeze(q_matrix,0) #need this because the autoencoder always expects batches of inputs!\n"," message=autoencoder.encode(q_matrix)[0]\n"," message_dict[task_index]=message\n","\n"," #state int-to-tuple dictionary\n"," state_int_to_tuple_dict={}\n"," for s in range(grid_dim**2):\n"," state_int_to_tuple_dict[s]=state_int_to_tuple(s)\n","\n"," #get state tensors\n"," state_tensors=get_state_tensors(1)\n","\n"," #create a dictionary containing all the student action probabilities (for all possible messages)\n"," action_probas_dict={}\n"," for message_index, message in message_dict.items():\n"," message=message_dict[message_index].to(device).detach()\n"," #create student Q-matrix and corresponding action probabilities\n"," Q=autoencoder.student(message, state_tensors)[0]\n"," action_probas=softy(Q)\n"," action_probas_dict[message_index]=action_probas\n","\n","\n"," #solving rates for the informed and misinformed students are saved in the corresponding arrays\n"," solving_rates = [np.zeros(len(label_dict)), np.zeros(len(label_dict))]\n"," solving_rates_worlds_avg=[[],[]]\n"," #iterate over gridworlds\n"," for world_index, wall_states in wall_state_dict.items():\n"," world_task_indices=[i for i,label in label_dict.items() if label[0]==world_index and label[2]!=student_init] #all tasks in this particular world\n"," world_quotas_info, world_quotas_random, world_quotas_misinfo =[], [], []\n"," G=graph_from_walls(wall_states)\n"," solving_rates_world = [np.zeros(len(world_task_indices)), np.zeros(len(world_task_indices)), np.zeros(len(world_task_indices))] #solving rates in this particular world\n"," #iterate over goal locations\n"," for j,task_index in enumerate(world_task_indices):\n"," goal_state=label_dict[task_index][2]\n"," #initialize environment\n"," env = SquareGridworld(student_init,goal_state,wall_states, lava)\n"," outcomes=env.get_outcomes()\n"," #dictionary to retreive next state and reward given current (s,a)\n"," next_states_dict={s:[outcomes[s,a][0] for a in [0,1,2,3]] for s in range(grid_dim**2)}\n"," #print(wall_states, student_init, goal_state)\n"," goal_distance=nx.dijkstra_path_length(G,student_init,goal_state, weight=None)\n"," #allowed steps are the optimum steps times some factor (if factor is 0.5 need to try both and average)\n"," if not step_rate_bool:\n"," max_steps_list=[round(solving_steps[task_index]-0.500001), round(solving_steps[task_index]+0.499999)]\n"," elif step_rate_bool:\n"," if (stepfactor*goal_distance)%1<0.01:\n"," max_steps_list=[round(stepfactor*goal_distance)]\n"," else:\n"," max_steps_list=[round(stepfactor*goal_distance-0.1), round(stepfactor*goal_distance+0.1)]\n","\n"," #iterate over informed (l=0) and misinformed (l=1) students\n"," for l in range(2):\n","\n"," rdrep_fraction=random_repeat if l==1 else 1 #informed student does not need several repetitions, as it gets one fixed message\n","\n"," for g in range(rdrep_fraction):\n"," if l==0 and g>0:\n"," continue\n","\n"," #random message for misinformed student, correct message for random and informed students\n"," message_index=task_index if l==0 else np.random.randint(0,len(message_dict))\n"," action_probas=action_probas_dict[message_index]\n","\n"," if method==\"no_learning\": #a stochastic method, actions have certain probabilities\n"," matrix_big=torch.zeros(size=(grid_dim**2,grid_dim**2), device=device) #create a big transition matrix\n"," for s in range(grid_dim**2):\n"," if s!=goal_state:\n"," for i,ns in enumerate(next_states_dict[s]):\n"," matrix_big[ns,s]+=action_probas[i,mt.floor(s/grid_dim),s%grid_dim] #need to have the \"+=\" here because there can be multiple identical transitions, e.g. for corner states\n"," else:\n"," matrix_big[s,s]=1\n","\n"," #iterate over the three students\n"," for max_steps in max_steps_list:\n"," max_steps=int(max_steps)\n","\n"," if method==\"no_learning\":\n"," #initialize state occupancy probabilities\n"," probas=0*probas\n"," probas[student_init]=1\n"," #For each student step, apply the transition matrix to the initial probability distribution to get the new probability distribution\n"," for rep in range(max_steps):\n"," probas=matrix_big@probas\n"," if not step_rate_bool:\n"," solving_rates_world[l][j]+=probas[goal_state]/rdrep_fraction*(1-abs(solving_steps[task_index]-max_steps))\n"," solving_rates[l][task_index]+=probas[goal_state]/rdrep_fraction*(1-abs(solving_steps[task_index]-max_steps))\n"," elif step_rate_bool:\n"," solving_rates_world[l][j]+=probas[goal_state]/(rdrep_fraction*len(max_steps_list))\n"," solving_rates[l][task_index]+=probas[goal_state]/(rdrep_fraction*len(max_steps_list))\n","\n","\n","\n"," elif method==\"simple_learning\": #a deterministic method (given the Q-matrix, we know exactly what the student's path will look like)\n"," state=student_init\n"," for step in range(max_steps):\n"," #take the action with maximum Q-value\n"," action=torch.argmax(action_probas[:,mt.floor(state/grid_dim),state%grid_dim])\n"," #shift its Q-value so that it is now lowest of all the four actions\n"," action_probas[action,mt.floor(state/grid_dim),state%grid_dim]=torch.min(action_probas[:,mt.floor(state/grid_dim),state%grid_dim])-1e-5\n"," state=next_states_dict[state][action]\n"," if state==goal_state:\n"," if not step_rate_bool:\n"," solving_rates_world[l][j]+=1/rdrep_fraction*(1-abs(solving_steps[task_index]-max_steps))\n"," solving_rates[l][task_index]+=1/rdrep_fraction*(1-abs(solving_steps[task_index]-max_steps))\n"," elif step_rate_bool:\n"," solving_rates_world[l][j]+=1/(rdrep_fraction*len(max_steps_list))\n"," solving_rates[l][task_index]+=1/(rdrep_fraction*len(max_steps_list))\n"," break\n","\n"," #print(f\"goal {goal_state}: solving rates info:{round(100*solving_rates[0][task_index])}%, misinfo:{round(100*solving_rates[1][task_index])}%\")\n"," for i in [0,1]:\n"," if len(solving_rates_world[i]>0): #assume that task numbers per world are approx. equal and then do average\n"," solving_rates_worlds_avg[i]+=[np.mean(solving_rates_world[i])]\n"," print(f\"in world {world_index} with walls {wall_states} and {len(world_task_indices)} tasks the solving rates are info:{round(100*np.mean(solving_rates_world[0]),2)}%, misinfo:{round(100*np.mean(solving_rates_world[1]),2)}%\")\n","\n"," print(\"We have the overall results:\")\n"," print(f\"Average solving rate for informed: ({round(100*np.mean(solving_rates[0]),2)} +/- {round(100*np.std(solving_rates[0])/mt.sqrt(len(solving_rates[0])),2)})%\")\n"," print(f\"Average solving rate for misinformed: ({round(100*np.mean(solving_rates[1]),2)} +/- {100*round(np.std(solving_rates[1])/mt.sqrt(len(solving_rates[1])),2)})%\")\n"," print(\"\")\n","\n"," #save the results\n"," if save_rates:\n"," if not os.path.exists(file_loc+\"student/\"+saving_folder):\n"," os.mkdir(file_loc+\"student/\"+saving_folder)\n"," if not step_rate_bool:\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_info_{method}_rate{rdwalker_base_rate}.txt\", solving_rates[0])\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_misinfo_{method}_rate{rdwalker_base_rate}.txt\", solving_rates[1])\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_worlds_info_{method}_rate{rdwalker_base_rate}.txt\", solving_rates_worlds_avg[0])\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_worlds_misinfo_{method}_rate{rdwalker_base_rate}.txt\", solving_rates_worlds_avg[1])\n"," elif step_rate_bool:\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_info_{method}_stepfactor{stepfactor}.txt\", solving_rates[0])\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_misinfo_{method}_stepfactor{stepfactor}.txt\", solving_rates[1])\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_worlds_info_{method}_stepfactor{stepfactor}.txt\", solving_rates_worlds_avg[0])\n"," np.savetxt(file_loc+\"student/\"+saving_folder+f\"/solving_rates_worlds_misinfo_{method}_stepfactor{stepfactor}.txt\", solving_rates_worlds_avg[1])\n","\n","\n","def random_walker_rates(saving_folder, language_number, step_rate_bool, rdwalker_base_rate, stepfactor, label_dict):\n"," '''\n"," Evaluates a language on the performance of the student by generating solving rates for all the possible tasks\n"," Two students are investigates:\n"," -random walker that takes completely random actions\n"," -smart random walker that takes completely random actions, but never runs into a wall\n"," ---\n"," INPUT\n"," saving_folder - String. Identifier of the folder where the solving rates should be saved\n"," language_number - integer. How many languages were generated for this set of trained worlds and trained tasks (to reduce variance)\n"," method - either \"no_learning\" or \"simple_learning\"\n"," stepfactor - integer. how many times the number of steps in the optimal solution do we give the student before the episode is ended\n"," step_rate_bool - boolean. which method do we do, either the stepfactor method (True) or the random walker rate method (False) (for determining how many steps are allowed)\n"," rdwalker_base_rate - float. we calculate how many steps the random walkers need to achieve the solving rate of rdwalker_base_rate for each task\n"," ---\n"," OUTPUT\n"," None. (Store data of solving rates for the two students in a corresponding file.\n"," '''\n","\n"," #state int-to-tuple dictionary\n"," state_int_to_tuple_dict={}\n"," probas=torch.zeros(grid_dim**2, device=device)\n"," for s in range(grid_dim**2):\n"," state_int_to_tuple_dict[s]=state_int_to_tuple(s)\n","\n"," #solving rates for the three students are saved in the corresponding arrays\n"," solving_steps_rd, solving_steps_smartrd = np.zeros(len(label_dict)), np.zeros(len(label_dict))\n"," solving_rates_rd, solving_rates_smartrd = np.zeros(len(label_dict)), np.zeros(len(label_dict))\n"," solving_ratesworlds_rd, solving_ratesworlds_smartrd = np.zeros(len(wall_state_dict)), np.zeros(len(wall_state_dict))\n","\n"," #iterate over gridworlds\n"," for world_index, wall_states in wall_state_dict.items():\n"," world_task_indices=[i for i,label in label_dict.items() if label[0]==world_index and label[2]!=student_init] #all tasks in this particular world\n"," solving_rates_rd_world, solving_rates_smartrd_world = np.zeros(len(world_task_indices)), np.zeros(len(world_task_indices))\n"," world_quotas_info, world_quotas_random, world_quotas_misinfo =[], [], []\n"," G=graph_from_walls(wall_states)\n"," #iterate over goal locations\n"," for j,task_index in enumerate(world_task_indices):\n"," goal_state=label_dict[task_index][2]\n"," #initialize environment\n"," env = SquareGridworld(student_init,goal_state,wall_states, lava)\n"," outcomes=env.get_outcomes()\n"," #dictionary to retreive next state and reward given current (s,a)\n"," next_states_dict={s:[outcomes[s,a][0] for a in [0,1,2,3]] for s in range(grid_dim**2)}\n"," goal_distance=nx.dijkstra_path_length(G,student_init,goal_state, weight=None)\n","\n"," #create matrices with random action probabilities\n"," matrix_big=torch.zeros(size=(grid_dim**2,grid_dim**2), device=device) #create a big transition matrix, each action has same likelihood\n"," matrix_big_nowall=torch.zeros(size=(grid_dim**2,grid_dim**2), device=device) #create a big transition matrix, wall actions are likelihood zero\n"," for s in range(grid_dim**2):\n"," if s!=goal_state and not (s in wall_states):\n"," next_states=next_states_dict[s]\n"," nowall_options=[ns for ns in next_states if ns!=s]\n"," for ns in next_states:\n"," matrix_big[ns,s]+=0.25 #need to have the \"+=\" here because there can be multiple identical transitions, e.g. for corner states\n"," for ns in nowall_options:\n"," matrix_big_nowall[ns,s]=1/len(nowall_options)\n","\n"," else:\n"," matrix_big[s,s]=1\n"," matrix_big_nowall[s,s]=1\n","\n"," #a) stepfactor method -> we give a stepfactor times shortest path length number of steps to solve the task.\n"," if step_rate_bool:\n"," #allowed steps are the optimum steps times some factor (if factor is 0.5 need to try both and average)\n"," if (stepfactor*goal_distance)%1<0.01:\n"," max_steps_list=[round(stepfactor*goal_distance)]\n"," else:\n"," max_steps_list=[round(stepfactor*goal_distance-0.1), round(stepfactor*goal_distance+0.1)]\n"," for max_steps in max_steps_list:\n"," #initialize state occupancy probabilities\n"," probas=0*probas\n"," probas[student_init]=1\n"," probas_nowall=torch.zeros(grid_dim**2, device=device)\n"," probas_nowall[student_init]=1\n"," for rep in range(max_steps):\n"," #1. regular random walker\n"," probas=matrix_big@probas\n"," #2. smart random walker\n"," probas_nowall=matrix_big_nowall@probas_nowall\n","\n"," solving_rates_rd[task_index]+=probas[goal_state]/len(max_steps_list)\n"," solving_rates_rd_world[j]+=probas[goal_state]/len(max_steps_list)\n"," solving_rates_smartrd[task_index]+=probas_nowall[goal_state]/len(max_steps_list)\n"," solving_rates_smartrd_world[j]+=probas_nowall[goal_state]/len(max_steps_list)\n","\n"," #b) random walker rate method -> we calculate the steps needed for the random walker to surpass a certain goal reaching probability\n"," if not step_rate_bool:\n"," #1.regular random walker\n"," prev_rate1, rate1, steps1=0,0,0\n"," while rate1 < rdwalker_base_rate:\n"," steps1+=1\n"," prev_rate1=rate1\n"," #initialize state occupancy probabilities\n"," probas=0*probas\n"," probas[student_init]=1\n"," #For each student step, apply the transition matrix to the initial probability distribution to get the new probability distribution\n"," for rep in range(steps1):\n"," probas=matrix_big@probas\n"," rate1=probas[goal_state]\n"," #linear interpolation to find out the approximate step number for reaching rdwalker_base_rate\n"," solving_steps_rd[task_index]=steps1-1+(rdwalker_base_rate-prev_rate1)/(rate1-prev_rate1)\n","\n"," #2.smart random walker\n"," prev_rate2, rate2, steps2=0,0,0 #for smart random walker\n"," while rate2 we don't analyze this gridworld\n"," if not nx.is_connected(G):\n"," continue\n"," for goal_state in range(grid_dim**2):\n"," if goal_state==student_init: #start location can not be a goal -> would be trivially solved\n"," continue\n"," elif goal_state in wall_states: #wall states can not be a goal -> they are not accessible for the agents\n"," continue\n"," else:\n"," label_dict[task_index]=[wall_index, student_init, goal_state]\n"," task_index+=1\n","\n","\n","\n","if qmat_gen:\n"," #first generate the perfect deterministic Q-matrices\n"," perfect_qdict=q_matrix_generator_deterministic(label_dict, wall_state_dict)\n"," '''\n"," lp = LineProfiler()\n"," lp_wrapper = lp(q_matrix_generator)\n"," lp_wrapper(wall_state_dict, perfect_q_matrices_dict)\n"," lp.print_stats()\n"," '''\n"," #Then use the DQN and the perfect matrices as a check\n"," q_matrix_dict=q_matrix_generator(label_dict, wall_state_dict,perfect_qdict, accuracy, max_attempts)\n"," #write calculated q matrices into file\n"," write_dict_into_pkl(q_matrix_dict,file_loc+\"teacher/q matrix dictionaries/\"+f\"q_matrices{qmat_save_code}.pkl\")\n"," write_dict_into_pkl(label_dict,file_loc+\"teacher/label dictionaries/\"+f\"q_matrices_labels{qmat_save_code}.pkl\")\n"," write_dict_into_pkl(wall_state_dict,file_loc+\"teacher/wall state dictionaries/\"+f\"wall_states{qmat_save_code}.pkl\")\n","\n","\n","if not qmat_gen:\n"," #read Q-matrices from a pre-generated file\n"," wall_state_dict=read_dict_from_pkl(file_loc+\"teacher/wall state dictionaries/\"+f\"wall_states{qmat_read_code}.pkl\")\n"," q_matrix_dict=read_dict_from_pkl(file_loc+\"teacher/q matrix dictionaries/\"+f\"q_matrices{qmat_read_code}.pkl\")\n"," label_dict=read_dict_from_pkl(file_loc+\"teacher/label dictionaries/\"+f\"q_matrices_labels{qmat_read_code}.pkl\")\n"," #transfer Q-matrices to correct device\n"," for key in q_matrix_dict.keys():\n"," q_matrix_dict[key] = q_matrix_dict[key].to(device)\n","\n","\n","#insert shortest path length from initial state to goal state into label dictionary (needed for language training with student feedback)\n","newlabel_dict={}\n","for i,label in label_dict.items():\n"," walls,init,goal=label\n"," wall_states=wall_state_dict[walls]\n"," G=graph_from_walls(wall_states)\n"," o=node_dist(G,init,goal)\n"," newlabel_dict[i]=[walls,init,goal,o]\n","label_dict=newlabel_dict\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":499,"status":"ok","timestamp":1702900618033,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"wl17kavC7gf_"},"outputs":[],"source":["#@title 18. EXECUTION - Language training (with or without student feedback)\n","\n","#autoencoder: specific properties/parameters\n","train_order=0 #0: no parameters frozen / 1:student parameters frozen / 2:autoencoder parameters frozen\n","batch_size=10 #batch the data (mini-batch gradient descent)\n","train_worlds=range(16) #these are the worlds we are training on (from the wall_state_dict)\n","#different indices describe different situations where we train on all mazes, but only a subset of the goals\n","train_goals_dict={0:[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15],1:[1,4,3,6,9,12,11,14],2:[2,5,8,7,10,13,15],3:[1,4,5,8,9,12,13],4:[2,3,6,7,10,11,14,15],5:[1,2,3,4,5,6,7],6:[8,9,10,11,12,13,14,15]}\n","train_goals_index=0 #which goal locations are used for training the language?\n","nonlinear_ae_train, nonlinear_std_train=True, True #do autoencoder/student networks have nonlinear activations in training?\n","\n","language_number_autoenc=1 #how many languages with the same parameters and different seeds are trained?\n","language_save_code=f\"test_language\" #this list and the two following need to have the same length (will be zipped together)\n","\n","if language_gen:\n","\n"," train_goals=train_goals_dict[train_goals_index] #pick the goals we train on\n"," for language_index in range(language_number_autoenc): #create several languages to reduce variance\n","\n"," language_save_codex=language_save_code\n"," if language_number_autoenc>0:\n"," language_save_codex=language_save_code+f\"_language{language_index}\"\n","\n"," goal_world_inds, _,_,_ = task_indices_sorter(label_dict, train_worlds, train_goals)\n"," #create the dataset for the network to train on from the training indices\n"," q_matrix_dict_train, label_dict_train = {},{}\n"," #filter out only the relevant tasks (maze in train_worlds, goal in train_goals)\n"," for task_counter,ind in enumerate(goal_world_inds):\n"," q_matrix_dict_train[task_counter]=q_matrix_dict[ind]\n"," label_dict_train[task_counter]=label_dict[ind]\n","\n"," #create dataset\n"," matrix_dataset = MatrixDataset(q_matrix_dict_train, label_dict_train)\n","\n"," #initialize network and optimizer\n"," autoencoder = ConvAutoEncoder(data_shape, K, nonlinear_ae_train, nonlinear_std_train).to(device)\n"," optimizer = torch.optim.Adam(autoencoder.parameters(), lr=learning_rate_autoenc)\n","\n"," '''\n"," #training without student feedback\n"," losses1, rec_losses1, spar_losses1 = train_autoencoder(autoencoder, optimizer, matrix_dataset, wall_state_dict,\n"," gamma_sparse,zeta_std, kappa, training_epochs, batch_size, False, train_order)\n","\n"," #save autoencoder parameters and losses\n"," torch.save(autoencoder.state_dict(), file_loc+\"autoencoder/autoencoder network parameters/\"+f\"params_autoenc{language_save_codex}.pt\")\n"," torch.save(optimizer.state_dict(), file_loc+\"autoencoder/optimizer parameters/\"+f\"optimizer{language_save_codex}.pt\")\n"," np.savetxt(file_loc+f\"autoencoder/losses/losses1_{language_save_codex}.txt\", losses1)\n"," np.savetxt(file_loc+f\"autoencoder/losses/rec_losses1_{language_save_codex}.txt\",rec_losses1)\n"," np.savetxt(file_loc+f\"autoencoder/losses/spar_losses1_{language_save_codex}.txt\",spar_losses1)\n","\n"," #re-initialize for new training with student feedback\n"," autoencoder = ConvAutoEncoder(data_shape, K, nonlinear_ae_train, nonlinear_std_train).to(device)\n"," optimizer = torch.optim.Adam(autoencoder.parameters(), lr=learning_rate_autoenc)\n"," '''\n","\n"," #training with student feedback\n"," losses2, rec_losses2, spar_losses2 = train_autoencoder(autoencoder, optimizer, matrix_dataset, wall_state_dict,\n"," gamma_sparse,zeta_std, kappa, training_epochs, batch_size, True, train_order)\n","\n"," goal_losses2=(losses2-(1-gamma_sparse)*rec_losses2-gamma_sparse*spar_losses2)/zeta_std\n","\n"," #save autoencoder parameters and losses\n"," torch.save(autoencoder.state_dict(), file_loc+\"autoencoder/autoencoder network parameters/\"+f\"params_autoenc{language_save_codex}.pt\")\n"," torch.save(optimizer.state_dict(), file_loc+\"autoencoder/optimizer parameters/\"+f\"optimizer{language_save_codex}.pt\")\n"," np.savetxt(file_loc+f\"autoencoder/losses/losses2_{language_save_codex}.txt\", losses2)\n"," np.savetxt(file_loc+f\"autoencoder/losses/rec_losses2_{language_save_codex}.txt\",rec_losses2)\n"," np.savetxt(file_loc+f\"autoencoder/losses/spar_losses2_{language_save_codex}.txt\",spar_losses2)\n"," np.savetxt(file_loc+f\"autoencoder/losses/goal_losses2_{language_save_codex}.txt\",goal_losses2)\n","\n","\n","#load stored autoencoder network parameters\n","if not language_gen:\n"," autoencoder = ConvAutoEncoder(data_shape, K, nonlinear_ae_plots, nonlinear_std_plots).to(device)\n"," autoencoder.load_state_dict(torch.load(file_loc+\"autoencoder/autoencoder network parameters/\"+f\"params_autoenc{language_code}.pt\", weights_only=True))\n"," autoencoder.eval()\n","\n"," '''\n"," #only load optimizer if necessary, e.g. if training was interrupted - then parameters can be specified\n"," optimizer=torch.optim.Adam(autoencoder.parameters(), lr=learning_rate_autoenc)\n"," optimizer.load_state_dict(torch.load(file_loc+\"autoencoder/optimizer parameters/\"+f\"optimizer{language_code}.pt\", map_location=torch.device(device), weights_only=True))\n"," '''\n","\n","#create a message dictionary, with indices corresponding to the task indices\n","message_dict={}\n","for task_index,q_matrix in q_matrix_dict.items():\n"," q_matrix=torch.unsqueeze(q_matrix,0) #need this because the autoencoder always expects batches of inputs!\n"," message=autoencoder.encode(q_matrix)[0]\n"," message_dict[task_index]=message\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","executionInfo":{"elapsed":5,"status":"ok","timestamp":1702900618626,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"8SKISQC6mP6s"},"outputs":[],"source":["#@title 19. EXECUTION - Calculate student performance (solving rates)\n","'''\n","test the students on how they process new information given a certain message\n","\n","methods: softmax with no learning and \"simple learning\", i.e. avoiding actions done before in this state\n","'''\n","\n","#Parameters for the evaluations\n","random_repeat=5 #randomness repetition for misinformed students\n","rdwalker_base_rates=[] #baseline random walker rates we set for each task (values like 0.1 for 10% and 0.2 for 20%)\n","stepfactors=[2] #number of steps we allow for the student in terms of shortest path length for each task\n","methods=[\"no_learning\"] #could choose \"no_learning\" (action choice by percentages) and/or \"simple_learning\" (greedy action choices, but if you come back to a state you take the next best action)\n","\n","#\"known\" for trained worlds or \"unknown\" for untrained worlds\n","if qmat_read_code==\"training4x4\":\n"," known_addon=\"known\"\n","elif qmat_read_code==\"test4x4\":\n"," known_addon=\"unknown\"\n","else:\n"," known_addon=None\n","\n","#Parameters specified by the languages we want to evaluate\n","goal_locs=[0] #list of trained goal locations\n","\n","language_nr_evaluation=1 #number of languages to evaluate (assume it's the same for all different goal locations)\n","language_codes=[f\"test_language\"]\n","saving_folders=[f\"test_language_{known_addon}\"]\n","evalute_nonlinear_ae, evaluate_nonlinear_std = True, True #are the activations in the autoencoder/student nonlinear?\n","\n","#first choice if we are doing 'closing the loop', i.e. using the student Q-matrices for creating messages (which are different for every language)\n","#and second choice if we take regular teacher Q-matrices for creating messages\n","#q_matrix_dict_list=[read_dict_from_pkl(file_loc+f\"closing the loop/test_language_language0/studentQmatrices.pkl\")]\n","q_matrix_dict_list=[q_matrix_dict]*len(language_codes)*language_nr_evaluation\n","\n","\n","\n","\n","gc.enable() #garbage collector enabled to free RAM\n","if student_evaluate:\n"," for step_rate_bool in [True,False]: #evaluation for all stepfactors (True) and for all random walker base rates (False) listed above\n"," if step_rate_bool:\n"," rdwalker_base_rate=0\n"," elif not step_rate_bool:\n"," stepfactor=0\n","\n"," #a)evaluate solving rates for different stepfactors\n"," if step_rate_bool:\n"," for stepfactor in stepfactors:\n"," i=0\n"," for lcode_eval, saving_folder in zip(language_codes, saving_folders):\n"," #first calculate random walker rates (identical for all languages, but need to do several times, because of saving procedure..)\n"," solving_steps_rd, solving_steps_smartrd=random_walker_rates(saving_folder, language_nr_evaluation, step_rate_bool, rdwalker_base_rate, stepfactor, label_dict)\n"," gc.collect()\n"," for language_nr in range(language_nr_evaluation):\n"," for method in methods:\n"," #next calculate rates for misinformed and informed students\n"," student_performance_evaluator(lcode_eval, saving_folder, language_nr, method, random_repeat, step_rate_bool, solving_steps_rd, rdwalker_base_rate, stepfactor, label_dict, q_matrix_dict_list[i], evalute_nonlinear_ae, evaluate_nonlinear_std)\n"," gc.collect()\n"," i+=1\n","\n"," #a)evaluate solving rates for different baseline random walker baseline solving rates\n"," elif not step_rate_bool:\n"," for rdwalker_base_rate in rdwalker_base_rates:\n"," i=0\n"," for lcode_eval, saving_folder in zip(language_codes, saving_folders):\n"," #first calculate random walker steps for the baseline solving rate (identical for all languages, but need to do several times, because of saving procedure..)\n"," solving_steps_rd, solving_steps_smartrd=random_walker_rates(saving_folder, language_nr_evaluation, step_rate_bool, rdwalker_base_rate, stepfactor, label_dict)\n"," gc.collect()\n"," for language_nr in range(language_nr_evaluation):\n"," for method in methods:\n"," #next calculate rates for misinformed and informed students and give them the steps the random walkers needed\n"," student_performance_evaluator(lcode_eval, saving_folder, language_nr, method, random_repeat, step_rate_bool, solving_steps_rd, rdwalker_base_rate, stepfactor, label_dict, q_matrix_dict_list[i], evalute_nonlinear_ae, evaluate_nonlinear_std)\n"," gc.collect()\n"," i+=1\n"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":1000},"executionInfo":{"elapsed":13163,"status":"error","timestamp":1702900632914,"user":{"displayName":"Tobias Wieczorek","userId":"09658587961490681575"},"user_tz":-60},"id":"iVZqhFg3_Vyg","outputId":"09fdab55-2fa3-46f3-cd12-88ed8be2303d"},"outputs":[],"source":["#@title 20. PLOTS - Message space plots (PCA and t-SNE and UMAP) (Figs. 2, S1, S2, S9-12)\n","'''\n","PCA and t-SNE dimensionality reductions of the messages\n","'''\n","\n","#general specs of the plots\n","world_number=len(plot_worlds)\n","\n","#parameters of t-sne plots\n","tsne_scaling=False #if True, then all dimensions will be of equal importance\n","tsne_iter=1500 #number of iterations for tsne algorithm\n","tsne_rdkernel=1\n","\n","#parameters of umap plots\n","umap_mindist=0.5 #minimum distance between two points (can prevent \"clumping\") -> (recommended: 0-1), standard 0.1\n","umap_rdkernel=1\n","\n","\n","#-> have to create the folder in Drive first\n","if not os.path.exists(file_loc+f\"message plots pca/{language_code}\"):\n"," os.mkdir(file_loc+f\"message plots pca/{language_code}\")\n","if not os.path.exists(file_loc+f\"message plots tsne/{language_code}\"):\n"," os.mkdir(file_loc+f\"message plots tsne/{language_code}\")\n","if not os.path.exists(file_loc+f\"message plots umap/{language_code}\"):\n"," os.mkdir(file_loc+f\"message plots umap/{language_code}\")\n","\n","#plot layout details\n","plt.rc('axes', labelsize=32)\n","plt.rc('xtick', labelsize=32)\n","plt.rc('ytick', labelsize=32)\n","plt.rc('axes', titlesize=32)\n","markersize_all, markersize_single=35,100 #marker size for the plots involving all words and just a single world respectively\n","plt.rcParams[\"axes.prop_cycle\"] = plt.cycler(\"color\", plt.cm.tab20(np.linspace(0,world_number/20,world_number))) #color cyclwer\n","cmap = ColorMap2DZiegler(range_x=(-0.5-(grid_dim/2-1), 0.5+(grid_dim/2-1)), range_y=(-0.5-(grid_dim/2-1), 0.5+(grid_dim/2-1))) #2d color map for goal locations\n","\n","#create a color dictionary mapping world index to color\n","world_color_dict={0:\"grey\"}\n","for world in range(1,grid_dim**2):\n"," xworld,yworld=world%grid_dim, mt.floor(world/grid_dim)\n"," cval=(grid_dim-1)/2\n"," world_color_dict[world]=cmap(xworld-cval,yworld-cval)/255\n","\n","\n","#transform messages from tensors to arrays so that PCA and t-SNE can be done on them - only keep the ones from relevant worlds\n","message_list=np.array([mssg.detach().cpu().numpy() for i,mssg in message_dict.items() if (label_dict[i][0] in plot_worlds and label_dict[i][2]!=student_init)])\n","label_list=[[i,j,k] for [i,j,k,l] in label_dict.values() if i in plot_worlds]\n","\n","print(f\"A total of {len(message_list)} messages are plotted\")\n","\n","#do the variance analysis of the message space\n","variance_analyzer(message_list, label_list, plot_worlds)\n","\n","#create variance-explained plot and goal location-/initialaction-colored plots of messages\n","pca_plotting(autoencoder, message_list, label_list, language_code, \"all\", cmap, markersize_all)\n","plt.show()\n","\n","#do PCA on all messages\n","pca=PCA(K)\n","message_pca=pca.fit_transform(message_list)\n","\n","\n","\n","#FIGURE WITH COLOR BY WALL POSITION IN 2D\n","fig,ax=plt.subplots()\n","ax.xaxis.set_tick_params(width=5, length=10)\n","ax.yaxis.set_tick_params(width=5, length=10)\n","#Scatterplot of the messages, PC1 vs PC2, color by world\n","for world in plot_worlds:\n"," world_indices=[i for i,label in enumerate(label_list) if label[0]==world]\n"," ax.scatter(message_pca[world_indices,0],message_pca[world_indices,1],label=f\"world {world}\", color=world_color_dict[world], s=markersize_all)\n","#make some adjustments to the axes so the legend is nicely placed and doesn't overlap with data points\n","plt.title(\"Color: wall position\", y=1.05)\n","plt.xlabel(f\"first PC\")\n","plt.ylabel(f\"second PC\")\n","xdatamin,xdatamax=min(message_pca[:,0]), max(message_pca[:,0])\n","ydatamin,ydatamax=min(message_pca[:,1]), max(message_pca[:,1])\n","xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n","ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n","xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n","#typical case of PCA\n","if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n","#for t-sne and umap have to take other scenario into account\n","else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n","plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n","if save_message_plots:\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA all messages, color world.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots pca/\"+language_code+f\"/PCA all messages, color world.png\", format=\"png\")\n","plt.show()\n","\n","\n","if do_tsne_message_plots:\n"," #do t-SNE on all messages\n"," if tsne_scaling:\n"," message_list = StandardScaler().fit_transform(message_list) #here we standardize data -> only needed if we want all axisensions to be of equal importance\n"," for tsne_perplexity in [10,50]:\n"," tsne = TSNE(perplexity=tsne_perplexity, max_iter=tsne_iter, random_state=tsne_rdkernel, init=\"pca\", learning_rate=\"auto\")\n"," message_tsne = tsne.fit_transform(message_list)\n"," #Now do T-SNE plots, first by goal location and initial action\n"," tsne_plotting(autoencoder, tsne_perplexity, tsne_iter, tsne_rdkernel, message_list, label_list, language_code, \"all\", tsne_scaling, cmap, markersize_all)\n"," plt.show()\n","\n"," #TSNE FIGURE WITH COLOR BY WALL POSITION IN 2D\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," for world in plot_worlds:\n"," world_indices=[i for i,label in enumerate(label_list) if label[0]==world]\n"," ax.scatter(message_tsne[world_indices,0],message_tsne[world_indices,1],label=f\"world {world}\", color=world_color_dict[world], s=markersize_all)\n"," #make some adjustments to the axes so the legend is nicely placed and doesn't overlap with data points\n"," plt.title(f\"Color: wall position, \"+r\"$\\pi$=\"+f\"{tsne_perplexity}\", y=1.05)\n"," xdatamin,xdatamax=min(message_tsne[:,0]), max(message_tsne[:,0])\n"," ydatamin,ydatamax=min(message_tsne[:,1]), max(message_tsne[:,1])\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n","\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," if save_message_plots:\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE all messages, color world, perplexity {tsne_perplexity}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots tsne/\"+language_code+f\"/t-SNE all messages, color world, perplexity {tsne_perplexity}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","if do_umap_message_plots:\n"," for umap_neighbors in [10,50]:\n"," #do UMAP on all messages\n"," umap_red = umap.UMAP(n_neighbors=umap_neighbors, min_dist=umap_mindist, random_state=umap_rdkernel)\n"," message_umap = umap_red.fit_transform(message_list)\n"," #Now do T-SNE plots, first by goal location and initial action\n"," umap_plotting(autoencoder, umap_neighbors, umap_mindist, umap_rdkernel, message_list, label_list, language_code, \"all\", cmap, markersize_all)\n"," plt.show()\n","\n"," #UMAP FIGURE WITH COLOR BY WALL POSITION IN 2D\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," for world in plot_worlds:\n"," world_indices=[i for i,label in enumerate(label_list) if label[0]==world]\n"," ax.scatter(message_umap[world_indices,0],message_umap[world_indices,1],label=f\"world {world}\", color=world_color_dict[world], s=markersize_all)\n"," #make some adjustments to the axes so the legend is nicely placed and doesn't overlap with data points\n"," plt.title(f\"Color: wall position, k={umap_neighbors}\", y=1.05)\n"," xdatamin,xdatamax=min(message_umap[:,0]), max(message_umap[:,0])\n"," ydatamin,ydatamax=min(message_umap[:,1]), max(message_umap[:,1])\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ymin, ymax = ydatamin - 0.07*(ydatamax-ydatamin), ydatamax + 0.07*(ydatamax-ydatamin)\n"," xmean, ymean=(xmax+xmin)/2, (ymax+ymin)/2\n","\n"," #typical case of PCA\n"," if xdatamax-xdatamin > ydatamax-ydatamin:\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/2)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," #for t-sne and umap have to take other scenario into account\n"," else:\n"," ax.set_xlim(xmean-1/2*(ymax-ymin), xmean+1*(ymax-ymin))\n"," ax.set_ylim(ymin, ymax)\n","\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," if save_message_plots:\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP all messages, color world, neighbors {umap_neighbors}.svg\", format=\"svg\")\n"," plt.savefig(file_loc+\"message plots umap/\"+language_code+f\"/UMAP all messages, color world, neighbors {umap_neighbors}.png\", format=\"png\")\n"," plt.show()\n","\n","\n","\n","\n","\n","\n","#Finally repeat the above plots for the single grid-worlds (here only need goal location and initial action plots)\n","for world in plot_worlds_single:\n"," print(f\"now looking at world {world}\")\n"," #transform data from tensors to arrays so that PCA can be done on them\n"," label_list_world=[j for i,j in label_dict.items() if j[0]==world]\n"," message_list_world=np.array([mssg.detach().cpu().numpy() for i,mssg in message_dict.items() if label_dict[i][0]==world])\n"," if len(message_list_world)>0:\n"," #create variance-explained plot and goal location-/initial action-colored plots of messages\n"," pca_plotting(autoencoder, message_list_world, label_list_world, language_code, world, cmap, markersize_single)\n"," if do_tsne_message_plots:\n"," for tsne_perplexity_singleworld in [2,5]:\n"," tsne_plotting(autoencoder, tsne_perplexity_singleworld, tsne_iter, tsne_rdkernel,message_list_world, label_list_world, language_code, world, tsne_scaling, cmap, markersize_single)\n"," if do_umap_message_plots:\n"," for umap_neighbors_singleworld in [2,5]:\n"," umap_plotting(autoencoder, umap_neighbors_singleworld, umap_mindist, umap_rdkernel, message_list_world, label_list_world, language_code, world, cmap, markersize_single)\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"dAG0VLrUgXeq"},"outputs":[],"source":["#@title 21. PLOTS - Autoencoder loss plots (Figs. 4a-d, S6)\n","\n","#plotting details\n","plt.rc('axes', labelsize=36)\n","plt.rc('xtick', labelsize=36)\n","plt.rc('ytick', labelsize=36)\n","plt.rc('axes', titlesize=36)\n","plt.rcParams['svg.fonttype']='none' #to make later editing of figures in inkscape easier\n","cmap=plt.get_cmap(\"bwr\")\n","\n","\n","\n","#load losses from case without student feedback\n","losses1=np.loadtxt(file_loc+f\"autoencoder/losses/nonlinear_losses1_factor1.txt\")\n","rec_losses1=np.loadtxt(file_loc+f\"autoencoder/losses/nonlinear_rec_losses1_factor1.txt\")\n","spar_losses1=np.loadtxt(file_loc+f\"autoencoder/losses/nonlinear_spar_losses1_factor1.txt\")\n","#load losses from case with student feedback\n","rec_losses2=np.loadtxt(file_loc+f\"autoencoder/losses/nonlinear_rec_losses2_zeta{zeta_lossplot}_factor1.txt\")\n","losses2=np.loadtxt(file_loc+f\"autoencoder/losses/nonlinear_losses2_zeta{zeta_lossplot}_factor1.txt\")\n","spar_losses2=np.loadtxt(file_loc+f\"autoencoder/losses/nonlinear_spar_losses2_zeta{zeta_lossplot}_factor1.txt\")\n","goal_losses2=np.loadtxt(file_loc+f\"autoencoder/losses/nonlinear_goal_losses2_zeta{zeta_lossplot}_factor1.txt\")\n","\n","\n","\n","#Plot reconstruction losses with and without student feedback\n","fig,ax = plt.subplots()\n","ax.xaxis.set_tick_params(width=5, length=10)\n","ax.yaxis.set_tick_params(width=5, length=10)\n","\n","ax.plot(range(epskip,len(losses1)), rec_losses1[epskip:len(losses1)], label=\"no feedback\", color=cmap(0.99), lw=3)\n","ax.plot(range(epskip,len(losses1)), rec_losses2[epskip:len(losses1)], label=\"with feedback\", color=cmap(0.15), lw=3)\n","\n","ax.set_xlabel(\"number of epochs\")\n","ax.set_ylabel(\"loss\")\n","plt.title(\"reconstruction loss\")\n","ax.set_yscale(\"log\")\n","ax.set_xticks([epskip,500,1000])\n","ax.set_yticks([50,100,200],[\"50\",\"100\",\"200\"])\n","plt.ylim(40,250)\n","if save_autoenc_lossplots:\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_reconstruction_loss_zeta{zeta_lossplot}_factor1.svg\",format=\"svg\")\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_reconstruction_loss_zeta{zeta_lossplot}_factor1.png\",format=\"png\")\n","\n","\n","\n","#Plot sparsity losses with and without student feedback\n","fig,ax = plt.subplots()\n","ax.xaxis.set_tick_params(width=5, length=10)\n","ax.yaxis.set_tick_params(width=5, length=10)\n","\n","ax.plot(range(epskip,len(losses1)), spar_losses1[epskip:len(losses1)], label=\"no feedback\", color=cmap(0.99), lw=3)\n","ax.plot(range(epskip,len(losses1)), spar_losses2[epskip:len(losses1)], label=\"with feedback\", color=cmap(0.15), lw=3)\n","\n","ax.set_xlabel(\"number of epochs\")\n","plt.title(\"sparsity loss\")\n","ax.set_yscale(\"log\")\n","ax.set_xticks([epskip,500,1000])\n","ax.set_yticks([20,50,100,200],[\"20\",\"50\",\"100\",\"200\"])\n","plt.ylim(20,200)\n","if save_autoenc_lossplots:\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_sparsity_loss_zeta{zeta_lossplot}_factor1.svg\",format=\"svg\")\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_sparsity_loss_zeta{zeta_lossplot}_factor1.png\",format=\"png\")\n","\n","\n","\n","#Plot SAE losses with and without student feedback plus goal finding losses\n","fig,ax = plt.subplots()\n","ax.xaxis.set_tick_params(width=5, length=10)\n","ax.yaxis.set_tick_params(width=5, length=10)\n","\n","ax.plot(range(epskip,len(losses1)), (1-gamma_sparse)*rec_losses1[epskip:len(losses1)]+gamma_sparse*spar_losses1[epskip:len(losses1)], label=\"no feedback\", color=cmap(0.99), lw=3)\n","ax.plot(range(epskip,len(losses1)), (1-gamma_sparse)*rec_losses2[epskip:len(losses1)]+gamma_sparse*spar_losses2[epskip:len(losses1)], label=\"with feedback\", color=cmap(0.15), lw=3)\n","ax.plot(range(epskip,len(losses1)), goal_losses2[epskip:len(losses1)], label=\"goal finding loss\", color=\"green\", lw=3)\n","\n","ax.set_xlabel(\"number of epochs\")\n","plt.title(\"SAE/goal finding losses\")\n","ax.set_yscale(\"log\")\n","ax.set_xticks([epskip,500,1000])\n","ax.set_yticks([1,10,100],[\"1\",\"10\",\"100\"])\n","plt.ylim(0.5,250)\n","if save_autoenc_lossplots:\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_combined_loss_zeta{zeta_lossplot}_factor1.svg\",format=\"svg\")\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_combined_loss_zeta{zeta_lossplot}_factor1.png\",format=\"png\")\n","\n","\n","\n","#Plot SAE loss difference between cases with and without student feedback\n","plt.figure()\n","plt.ticklabel_format(style='plain')\n","axt = plt.gca()\n","axt.tick_params(width=5, length=10)\n","\n","plt.plot(range(epskip,len(losses1)), (1-gamma_sparse)*rec_losses1[epskip:len(losses1)]+gamma_sparse*spar_losses1[epskip:len(losses1)]-((1-gamma_sparse)*rec_losses2[epskip:len(losses1)]+gamma_sparse*spar_losses2[epskip:len(losses1)]), label=\"loss difference\", color=\"black\", lw=3)\n","plt.plot(range(epskip,len(losses1)), np.zeros(len(losses1))[epskip:] , color=\"gray\", linestyle=\"dashed\",lw=3) #mark the 0 with a dashed line\n","\n","plt.xlabel(\"number of epochs\")\n","plt.xticks([epskip,500,1000])\n","plt.title(\"SAE loss difference\")\n","plt.ylim(-65,5)\n","if save_autoenc_lossplots:\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_lossdifference_zeta{zeta_lossplot}_factor1.svg\",format=\"svg\")\n"," plt.savefig(file_loc+f\"autoencoder/nonlinear_lossdifference_zeta{zeta_lossplot}_factor1.png\",format=\"png\")\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"BbtHsWD_-3f4"},"outputs":[],"source":["#@title 22. PLOTS - Student performance plots (different goal locations trained), (Figs, 4e-h, 5b-e, S3-5)\n","\n","\n","#label dictionaries for all trained and unknown mazes respectively\n","label_dict_train=read_dict_from_pkl(file_loc+\"teacher/label dictionaries/\"+f\"q_matrices_labelstraining_4x4.pkl\")\n","label_dict_test=read_dict_from_pkl(file_loc+\"teacher/label dictionaries/\"+f\"q_matrices_labelstest_4x4.pkl\")\n","\n","#plot settings\n","plt.rc('axes', labelsize=36)\n","plt.rc('xtick', labelsize=36)\n","plt.rc('ytick', labelsize=36)\n","plt.rc('axes', titlesize=36)\n","cmap=plt.get_cmap(\"bwr\")\n","\n","#a function to put p-value stars over bars\n","def add_pvalues(rects_x, rects_y, rect_width, pvalues, critical_p=0.05):\n"," '''\n"," rects_x - list of x-positions of the bars\n"," rects_y - list of bar heights (plus error bar)\n"," rect_width - rectangle width\n"," pvalues - list of pvalues (in order p01,p02,...,p0n,p12,...p1n,...,p n-1 n)\n"," critical_p - the critical p-value (if we are below that, we put a star in the plot)\n"," '''\n"," line_y = []\n"," y_offset_init=3\n"," y_offset_between=8\n"," ycounter=0 #counter for how many star bars we have already placed\n"," pcounter=0\n"," for i, [rect1_x, rect1_y] in enumerate(zip(rects_x, rects_y)):\n"," for j, [rect2_x, rect2_y] in enumerate(zip(rects_x, rects_y)):\n"," if i < j:\n"," pvalue = pvalues[pcounter]\n"," pcounter+=1\n"," star = ''\n"," if pvalue < critical_p:\n"," star='*'\n"," if star:\n"," x = (rect1_x + rect2_x) / 2\n"," y = max(rect1_y,rect2_y) + y_offset_init+ycounter*y_offset_between\n"," line_y.append(y)\n"," ax.annotate(star, xy=(x, y), fontsize=20,ha='center')\n"," ax.plot([rect1_x , rect2_x],[y, y], lw=3, c='black')\n"," ycounter+=1\n","\n","\n","\n","for method in [\"no_learning\"]: #iterate over learning methods\n"," method_add_on=\"simple learning\" if method==\"simple_learning\" else \"no learning\"\n"," for solving_rate_method in [\"stepfactor\"]:\n"," param= stepfactor_goalloc_plots if solving_rate_method==\"stepfactor\" else rdrate_goalloc_plots\n"," #potentially chuck out languages if the performances are bad\n"," chucked_indices=[[] for z in goal_groups_plots]\n"," for plot in [0,1,2,3]: #iterate over performance on known and unknown tasks\n","\n"," goal_groups_adjusted=list(goal_groups_plots)\n"," if plot in [1,3] and 0 in goal_groups_plots:\n"," goal_groups_adjusted.remove(0)\n"," goal_groups_adjusted=np.array(goal_groups_adjusted)\n","\n"," ldict=label_dict_train if plot in [0,1] else label_dict_test\n"," known_addon=\"trained\" if plot in [0,1] else \"unknown\"\n"," avg_rates=[[],[],[],[]] if solving_rate_method==\"stepfactor\" else [[],[]] #average solving rates over all tasks\n"," sem_rates=[[],[],[],[]] if solving_rate_method==\"stepfactor\" else [[],[]] #standard error of the mean in the average solving rates over all tasks\n"," #for t-tests\n"," p01_list,p02_list,p03_list,p12_list,p13_list,p23_list=[],[],[],[],[],[]\n"," plists=[p01_list,p02_list,p03_list,p12_list,p13_list,p23_list]\n"," for m in goal_groups_adjusted:\n","\n"," folder=folders_goalloc_plots[m]+\"_known\" if plot in [0,1] else folders_goalloc_plots[m]+\"_unknown\"\n","\n"," rates=[[],[],[],[]] if solving_rate_method==\"stepfactor\" else [[],[]] #solving rates from all tasks listed for the three students (random walker is separate)\n","\n"," for language in range(language_nr_goalloc_plots):\n"," single_rates=[[],[],[],[]] if solving_rate_method==\"stepfactor\" else [[],[]]\n"," #load the solving rates\n"," for k,student in enumerate([\"info\",\"misinfo\"]):\n"," single_rates[k]=np.loadtxt(file_loc+\"student/\"+f\"{folder}/{folder}_language{language}/\"+f\"solving_rates_{student}_{method}_{solving_rate_method}{param}.txt\")\n"," if solving_rate_method==\"stepfactor\":\n"," single_rates[2]=np.loadtxt(file_loc+\"student/\"+f\"{folder}/{folder}_language{language}/\"+f\"solving_rates_rdwalkersmart_{solving_rate_method}{param}.txt\")\n"," single_rates[3]=np.loadtxt(file_loc+\"student/\"+f\"{folder}/{folder}_language{language}/\"+f\"solving_rates_rdwalker_{solving_rate_method}{param}.txt\")\n","\n"," #filter out only the solving rates from tasks with trained goal locations\n"," if plot==0 or plot==2:\n"," for std in range(len(single_rates)):\n"," newrates=[]\n"," i=0\n"," for j,label in ldict.items():\n"," if label[2]!=student_init:\n"," if label[2] in train_goals_dict[m]:\n"," newrates+=[single_rates[std][i]]\n"," i+=1\n"," single_rates[std]=newrates\n","\n"," #filter out only the solving rates from tasks with unknown goal locations\n"," if plot==1 or plot==3:\n"," for std in range(len(single_rates)):\n"," newrates=[]\n"," i=0\n"," for j,label in ldict.items():\n"," if label[2]!=student_init:\n"," if not(label[2] in train_goals_dict[m]):\n"," newrates+=[single_rates[std][i]]\n"," i+=1\n"," single_rates[std]=newrates\n","\n"," #chuck the language out if the informed student is worse than the misinformed student!\n"," if chuck_out:\n"," if (plot==0 and (np.mean(single_rates[0])avg_rates[1][t] else 1\n"," p02val=p02_list[t] if avg_rates[0][t]>avg_rates[2][t] else 1\n"," p03val=p03_list[t] if avg_rates[0][t]>avg_rates[3][t] else 1\n"," p12val=p12_list[t] if avg_rates[1][t]>avg_rates[2][t] else 1\n"," p13val=p13_list[t] if avg_rates[1][t]>avg_rates[3][t] else 1\n"," p_values=[p01val, p02val,1,1,1,1]\n"," add_pvalues(rects_x,rects_y,rect_width,p_values)\n","\n","\n"," ax.set_ylim(0,120)\n"," ax.set_ylabel(\"tasks solved (%)\")\n"," if plot==0:\n"," plt.title(f\"student evaluated at trained goals for trained mazes (0 and 1 wall states)\")\n"," elif plot==1:\n"," plt.title(f\"student evaluated at unknown goals for trained mazes (0 and 1 wall states)\")\n"," elif plot==2:\n"," plt.title(f\"student evaluated at trained goals for unknown mazes (2 wall states)\")\n"," elif plot==3:\n"," plt.title(f\"student evaluated at unknown goals for unknown mazes (2 wall states)\")\n"," ax.set_xlim(-0.25,len(goal_groups_plots)-0.25)\n"," ax.set_yticks([0,20,40,60,80,100])\n","\n","\n"," ax.set_xlabel(r\"goal locations trained\")\n"," #plt.legend(bbox_to_anchor=(1,1), loc=\"upper left\", fontsize=25) #put legend outside the plot so that no lines are covered!\n"," if save_goalloc_plots:\n"," plt.savefig(file_loc+\"student/\"+f\"{folders_goalloc_plots[0]}\"+f\"/solvingrates_{method}_{solving_rate_method}_zeta5_factor1_plot{plot}.png\", bbox_inches='tight', format=\"png\")\n"," plt.savefig(file_loc+\"student/\"+f\"{folders_goalloc_plots[0]}\"+f\"/solvingrates_{method}_{solving_rate_method}_zeta5_factor1_plot{plot}.svg\", format=\"svg\")\n"," plt.show()\n","\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"VIjmfdzZXNMa"},"outputs":[],"source":["#@title 23. PLOTS - Closing the loop (Figs. 5a, S7)\n","\n","cmap = ColorMap2DZiegler(range_x=(-0.5-(grid_dim/2-1), 0.5+(grid_dim/2-1)), range_y=(-0.5-(grid_dim/2-1), 0.5+(grid_dim/2-1))) #2d color map for goal locations\n","cmap2=plt.get_cmap(\"bwr\")\n","softy=nn.Softmax(dim=0)\n","label_list=[[i,j,k] for [i,j,k,l] in label_dict.values()]\n","pca1, pca2=PCA(grid_dim**2*4), PCA(K) #q matrix PCA and message PCA\n","cval=(grid_dim-1)/2 #central value for reference if we use color by goal location or wall position\n","state_tensors=get_state_tensors(1) #for autoencoder forward function\n","\n","#create color dictionary for goal and wall state positions\n","color_dict={0:\"grey\"}\n","for s in range(1,grid_dim**2):\n"," #Scatterplot of the q-matrices, PC1 vs PC2\n"," x,y=s%grid_dim, mt.floor(s/grid_dim)\n"," color_dict[s]=cmap(x-cval,y-cval)/255\n","\n","#plot settings\n","plt.rc('axes', labelsize=32)\n","plt.rc('xtick', labelsize=32)\n","plt.rc('ytick', labelsize=32)\n","plt.rc('axes', titlesize=32)\n","plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n","\n","\n","print(\"\")\n","print(f\"Analyzing the language {language_code_closingloop}\")\n","print(\"\")\n","if not os.path.exists(file_loc+f\"closing the loop/{language_code_closingloop}\"):\n"," os.mkdir(file_loc+f\"closing the loop/{language_code_closingloop}\")\n","\n","#load the language\n","autoencoder = ConvAutoEncoder(data_shape, K, closingloop_nonlinear_ae, closingloop_nonlinear_std).to(device)\n","autoencoder.load_state_dict(torch.load(file_loc+\"autoencoder/autoencoder network parameters/\"+f\"params_autoenc{language_code_closingloop}.pt\", weights_only=True))\n","autoencoder.eval()\n","\n","#create the (regular) messages\n","message_dict={}\n","for task_index,q_matrix in enumerate(q_matrix_dict.values()):\n"," q_matrix=torch.unsqueeze(q_matrix,0) #need this because the autoencoder always expects batches of inputs!\n"," message=autoencoder.encode(q_matrix)[0]\n"," message_dict[task_index]=message\n","\n","#Data analysis\n","message_list=[mssg.detach().cpu().numpy() for i,mssg in message_dict.items()]\n","#1.PCA of the original Q-matrices\n","q_matrix_list=np.array([torch.flatten(Q).detach().numpy() for Q in q_matrix_dict.values()])\n","#2.PCA of the student Q-matrices\n","q_matrixstd_list=np.array([torch.flatten(autoencoder.student(torch.tensor([m]),state_tensors)[0]).detach().numpy() for m in message_list])\n","qmatrix_dict_studentq={k:autoencoder.student(torch.tensor([m]),state_tensors)[0] for k,m in enumerate(message_list)}\n","write_dict_into_pkl(qmatrix_dict_studentq,file_loc+f\"closing the loop/{language_code_closingloop}/studentQmatrices.pkl\") #save the Q-matrices belonging to this language\n","#3.PCA of the student probability matrices\n","p_matrixstd_list=np.array([torch.flatten(softy(autoencoder.student(torch.tensor([m]),state_tensors)[0])).detach().numpy() for m in message_list])\n","qmatrix_dict_studentp={k:softy(autoencoder.student(torch.tensor([m]),state_tensors)[0]) for k,m in enumerate(message_list)}\n","#4.PCA of messages from student Q-matrices\n","mstd_fromq_list=np.array([autoencoder.encode(torch.unsqueeze(autoencoder.student(torch.tensor([m]),state_tensors)[0],0))[0].detach().numpy() for m in message_list])\n","message_dict_studentq={k:m for k,m in enumerate(mstd_fromq_list)}\n","#5.PCA of messages from student probability matrices\n","mstd_fromp_list=np.array([autoencoder.encode(torch.unsqueeze(softy(autoencoder.student(torch.tensor([m]),state_tensors)[0]),0))[0].detach().numpy() for m in message_list])\n","message_dict_studentp={k:m for k,m in enumerate(mstd_fromp_list)}\n","\n","\n","datalist=[q_matrix_list, q_matrixstd_list, p_matrixstd_list, mstd_fromq_list, mstd_fromp_list]\n","savelist=[\"teacherQ\", \"studentQ\", \"studentP\", \"message_studentQ\", \"message_studentP\"]\n","printlist=[\"teacher Q-matrices\", \"student Q-matrices\", \"student P-matrices\", \"messages from student Q-matrices\", \"messages from student P-matrices\"]\n","for t, [data, savedata,printdata] in enumerate(zip(datalist,savelist,printlist)):\n"," print(f\"Doing PCA of {printdata}\")\n"," print(\"\")\n","\n"," #first do variance analysis\n"," variance_analyzer(data, label_list, plot_worlds)\n","\n"," pca=pca1 if t<3 else pca2 #different dimensionalities for Q-matrix space and message space require different PCAs\n"," data_pca=pca.fit_transform(data)\n","\n"," #PLOT VARIANCE EXPLAINED BY PC COMPONENT\n"," fig,ax = plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," var_arr=np.round(pca.explained_variance_ratio_*100,2)\n"," ax.set_ylim(0,100)\n"," if pca==pca1:\n"," plt.bar(range(11)[1:],var_arr[:10], color=\"gray\") #for each message dimension have one PC\n"," plt.xticks(range(11)[1:])\n"," else:\n"," plt.bar(range(K+1)[1:],var_arr[:K], color=\"gray\") #for each message dimension have one PC\n"," plt.xticks(range(K+1)[1:])\n","\n"," plt.xlabel(f\"PC index\")\n"," plt.ylabel(r\"Variance\" \"\\n\" r\"explained (%)\")\n"," if save_closingloop_plots:\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_varexplained.png\", format=\"png\")\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_varexplained.svg\", format=\"svg\")\n","\n","\n"," #PCA PLOT - COLOR BY WORLD (WALL PLACEMENT)\n"," fig,ax = plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," for world in range(16):\n"," world_indices=[k for k,label in enumerate(label_list) if label[0]==world]\n"," ax.scatter(data_pca[world_indices,0],data_pca[world_indices,1],label=f\"world {world}\", color=color_dict[world])\n"," ax.set_xlabel(\"first PC\")\n"," ax.set_ylabel(\"second PC\")\n"," xdatamin,xdatamax=min(data_pca[:,0]), max(data_pca[:,0])\n"," ydatamin,ydatamax=min(data_pca[:,1]), max(data_pca[:,1])\n"," ymean=(ydatamax+ydatamin)/2\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/3)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," plt.title(\"Color: wall position\")\n"," if save_closingloop_plots:\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_colorworld.png\", format=\"png\")\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_colorworld.svg\", format=\"svg\")\n","\n","\n"," #PCA PLOT - COLOR BY GOAL LOCATION\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," for goal in range(16):\n"," goal_indices=[k for k,label in enumerate(label_list) if label[2]==goal]\n"," ax.scatter(data_pca[goal_indices,0],data_pca[goal_indices,1],label=f\"goal {goal}\", color=color_dict[goal])\n"," ax.set_xlabel(\"first PC\")\n"," ax.set_ylabel(\"second PC\")\n"," xdatamin,xdatamax=min(data_pca[:,0]), max(data_pca[:,0])\n"," ydatamin,ydatamax=min(data_pca[:,1]), max(data_pca[:,1])\n"," ymean=(ydatamax+ydatamin)/2\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/3)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," plt.title(\"Color: goal location\")\n"," if save_closingloop_plots:\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_colorgoal.png\", format=\"png\")\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_colorgoal.svg\", format=\"svg\")\n","\n","\n"," #PCA PLOT - COLOR BY INITIAL ACTION\n"," fig,ax=plt.subplots()\n"," ax.xaxis.set_tick_params(width=5, length=10)\n"," ax.yaxis.set_tick_params(width=5, length=10)\n"," softy=nn.Softmax(dim=0)\n"," for l,message in enumerate(message_list):\n"," state_tensors=get_state_tensors(1)\n"," Q=autoencoder.student(torch.tensor([message]), state_tensors)[0]\n"," probas=softy(Q[:,0%grid_dim,mt.floor(0/grid_dim)]).detach().cpu().numpy()\n"," #combine the four probabilities into a 2-dimensional vector using the four corners of a square with four different colours\n"," extr=0.5+grid_dim/2-1\n"," comb2d=probas[0]*np.array([extr,extr])+probas[1]*np.array([extr,-extr])+probas[2]*np.array([-extr,-extr])+probas[3]*np.array([-extr,extr])\n"," ax.scatter(data_pca[l,0],data_pca[l,1],color=cmap(comb2d[0],comb2d[1])/255)\n","\n"," ax.set_xlabel(\"first PC\")\n"," ax.set_ylabel(\"second PC\")\n"," xdatamin,xdatamax=min(data_pca[:,0]), max(data_pca[:,0])\n"," ydatamin,ydatamax=min(data_pca[:,1]), max(data_pca[:,1])\n"," ymean=(ydatamax+ydatamin)/2\n"," xmin, xmax = xdatamin - 0.07*(xdatamax-xdatamin), xdatamax + 0.07*(xdatamax-xdatamin) #x axis limits a bit larger than data range to not cut points in the middle\n"," ax.set_xlim(xmin,xmax+(xmax-xmin)/3)\n"," ax.set_ylim(ymean-1/2*(xmax-xmin), ymean+1/2*(xmax-xmin)) #manually set equal aspect ratio, it was the only way that worked..\n"," plt.title(\"Color: first student action\")\n"," if save_closingloop_plots:\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_coloraction.png\", format=\"png\")\n"," plt.savefig(file_loc+f\"closing the loop/{language_code_closingloop}/{savedata}_coloraction.svg\", format=\"svg\")\n","\n"," plt.show()\n"," print(\"\\n\"*3)\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"CN-D5g7q1Bcf"},"outputs":[],"source":["#@title 24. EXECUTION - create task vectors and do topographical similarity analysis\n","'''\n","#For each task, create a goal/wall indicator vector as follows (dimensionality four):\n","#-the first two entries are the goal state's x- and y coordinates\n","#-the last two entries are the wall state's x- and y coordinates\n","#-the distance between two such vectors is simply the Euclidean distance\n","#-we create weightings for goal and wall by simply multiplying the first and last two components (respectively) with a number corresponding to that weighting\n","\n","#-the case of no wall state is covered by artificially placing the wall state at the top right corner outside of the maze\n","#(this way the distances are most meaningful, as this is symmetric w.r.t. action choice up vs. right and also a natural continuation of walls becoming less\n","#important overall if they are further away from the initial location)\n","'''\n","\n","#create dictionary mapping state numbers to coordinates as such:\n","#here the no-wall-state situation maps the wall state to s=(4,4)\n","#################\n","#3,0 3,1 3,2 3,3#\n","#2,0 2,1 2,2 2,3#\n","#1,0 1,1 1,2 1,3#\n","#0,0 0,1 0,2 0,3#\n","#################\n","w_goal=goal_reward\n","w_wall=wall_reward\n","#while the x- and y-axes start in the bottom left corner\n","state_coord_dict={i:[mt.floor(i/grid_dim), i%grid_dim] for i in range(16)}\n","state_coord_dict[0]=[4,4] #the goal is never at zero, so no issue here\n","indicatorvector_dict={}\n","indicatorvector_dictw={}\n","for task, [wall_index, init_state, goal_state, goalsplength] in label_dict.items():\n"," #get coordinates for wall state and goal state, append them\n"," wall_coord=np.array(state_coord_dict[wall_index])\n"," goal_coord=np.array(state_coord_dict[goal_state])\n"," overall_coord=np.append(goal_coord,wall_coord)\n"," overall_coordw=np.append(w_goal*goal_coord,w_wall*wall_coord)\n"," #put entry into dictionary\n"," indicatorvector_dict[task]=overall_coord\n"," indicatorvector_dictw[task]=overall_coordw\n","\n","\n","'''\n","#We calculate distances between all message pairs in a language and corresponding distances between a related quantity\n","#So far this quantity are either the ground truth Q-matrices or the ground truth action probability matrices\n","#In a next step, we then bin the meaning pairs (Q matrix or action proba matrix pairs) according to their distance in bins\n","#with equal number of entries and show the corresponding distance average and standard deviation in the signal pairs (message pairs)\n","\n","#currently: 25200 pairs of signals and meanings are compared\n","'''\n","#flatten and transfer to numpy\n","qdict={i:torch.flatten(q).detach().numpy() for i,q in q_matrix_dict.items()}\n","vdict=indicatorvector_dict\n","vwdict=indicatorvector_dictw\n","\n","#pairs distances calculation for ground truth Q-matrices and ground truth action probability matrices (see previous block)\n","qdist=np.array([np.linalg.norm(qdict[i]-qdict[j], norm_topo) for i in range(225) for j in range(i+1,225)])\n","vwdist=np.array([np.linalg.norm(vwdict[i]-vwdict[j], norm_topo) for i in range(225) for j in range(i+1,225)])\n","\n","#sort the meaning distance pairs by distance (Q- and P-matrices) and store the sorted keys\n","qdist_dict={i:q for i,q in enumerate(qdist)}\n","vwdist_dict={i:vw for i,vw in enumerate(vwdist)}\n","qkeylist=list(dict(sorted(qdist_dict.items(), key=lambda item: item[1])).keys())\n","vwkeylist=list(dict(sorted(vwdist_dict.items(), key=lambda item: item[1])).keys())\n","\n","for lcode_topo in [f\"nonlinear_nostudent_language{i}\" for i in range(5)]+[f\"nonlinear_goallocs0_zeta5_language{i}\" for i in range(5)]:\n"," #PREPARATION: first create the message dictionary corresponding to a certain language\n"," #load stored autoencoder network parameters\n"," autoencoder_topo = ConvAutoEncoder(data_shape, K, nonlinear_ae_plots, nonlinear_std_plots).to(device)\n"," autoencoder_topo.load_state_dict(torch.load(file_loc+\"autoencoder/autoencoder network parameters/\"+f\"params_autoenc{lcode_topo}.pt\", weights_only=True))\n"," autoencoder_topo.eval()\n"," #create a message dictionary, with indices corresponding to the task indices\n"," message_dict2={}\n"," for task_index,q_matrix in q_matrix_dict.items():\n"," q_matrix=torch.unsqueeze(q_matrix,0) #need this because the autoencoder always expects batches of inputs!\n"," message=autoencoder_topo.encode(q_matrix)[0]\n"," message_dict2[task_index]=torch.flatten(message).detach().numpy()\n","\n"," #TOPOGRAHIC SIMILARITY ANALYSIS\n"," #pairs distances calculation for messages\n"," mdist=np.array([np.linalg.norm(message_dict2[i]-message_dict2[j]) for i in range(225) for j in range(i+1,225)])\n","\n"," el_per_bin=int(len(qdist)/n_bins_topo) #divisibility required!!\n"," #for every language, for every bin store average meaning value and corresponding average signal value\n"," xqvals, yqvals, xpvals, ypvals, xvvals, yvvals, xvwvals, yvwvals=[], [], [], [], [], [], [], []\n"," for k in range(n_bins_topo): #again divisibility of number of pairs by bin number is needed\n"," qkeylist_chunk=qkeylist[el_per_bin*k:el_per_bin*(k+1)]\n"," vwkeylist_chunk=vwkeylist[el_per_bin*k:el_per_bin*(k+1)]\n","\n"," mqchunk=mdist[qkeylist_chunk]\n"," mvwchunk=mdist[vwkeylist_chunk]\n"," qchunk=qdist[qkeylist_chunk]\n"," vwchunk=vwdist[vwkeylist_chunk]\n","\n"," xqvals+=[np.average(qchunk)]\n"," xvwvals+=[np.average(vwchunk)]\n"," yqvals+=[np.average(mqchunk)]\n"," yvwvals+=[np.average(mvwchunk)]\n","\n","\n"," np.savetxt(file_loc+ f\"topographic similarity/q vs m {n_bins_topo} bins {norm_topo}-norm {lcode_topo}\",np.array([xqvals,yqvals]))\n"," np.savetxt(file_loc+ f\"topographic similarity/vw vs m {n_bins_topo} bins {norm_topo}-norm {lcode_topo}\",np.array([xvwvals,yvwvals]))\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"FwyEDysG_kss"},"outputs":[],"source":["#@title 25. PLOTS - Topographical similarity analysis (Figs. 3b, c)\n","'''\n","#See previous block - here are the plots\n","#basically what we see is how far are the messages from each other, when the ground truth meanings\n","#(i.e. Q-matrices or proba matrices) have a certain distance from each other\n","\n","#as every language is different, we also normalize by average message pair distance in the second plot\n","#and in both cases we compare 5 languages each\n","\n","#left column: not normalized, right column: normalized by average message pair distance\n","#first row: Q-matrices\n","#second row: P-matrices\n","'''\n","\n","fig, ax =plt.subplots(1,2, figsize=(14,5.6))\n","cmap=plt.get_cmap(\"bwr\")\n","plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n","plt.rc('axes', labelsize=28)\n","plt.rc('xtick', labelsize=28)\n","plt.rc('ytick', labelsize=28)\n","\n","for i, (language_prefix, color, color2, label) in enumerate(zip([\"nonlinear_nostudent\", \"nonlinear_goallocs0_zeta5\"], [cmap(0.99), cmap(0.15)], [cmap(0.75), cmap(0.25)], [\"no feedb.\", \"with feedb.\"])):\n"," for j, meaning in enumerate([\"q\",\"vw\"]):\n","\n"," yvals=np.zeros((5,n_bins_topo))\n","\n"," for l,lcode_topo in enumerate([f\"{language_prefix}_language{k}\" for k in range(5)]):\n"," x, y=np.loadtxt(file_loc+ f\"topographic similarity/{meaning} vs m {n_bins_topo} bins {norm_topo}-norm {lcode_topo}\")\n"," if l==0:\n"," xvals=x\n"," yvals[l]=y\n","\n"," #average across languages\n"," yvals, yvals_std =np.mean(yvals, axis=0), np.std(yvals, axis=0)\n","\n"," #linear curve fits for both scenarios\n"," popt1, pcov1=curve_fit(linfunc, xvals, yvals, sigma=yvals_std, absolute_sigma=True)\n","\n"," #parameter formatting\n"," m1, b1=np.round(popt1[0],3), np.round(popt1[1],3)\n"," m1_std, b1_std=np.round(np.sqrt(np.diagonal(pcov1))[0],3), np.round(np.sqrt(np.diagonal(pcov1))[1],3)\n","\n"," #first plot: regular message distances (averaged over languages)\n"," ax[j].errorbar(xvals, yvals, yerr=yvals_std, c=color, fmt=\"o\", ms=5, label=label+f\": m={m1}\"+u\"\\u00B1\"+f\"{m1_std}\", lw=3)\n"," ax[j].plot(xvals, linfunc(xvals, *popt1), color=color, lw=3)\n"," ax[j].legend(loc=2, fontsize=18)\n","\n","\n","xlims=[[0,6.5], [0,9]] if norm_topo==2 else [[0,35],[0,15]] if norm_topo==1 else None\n","ylims =[0,1.2]\n","ax[0].set_xlabel(f\"teacher Q-matrix distance\")\n","ax[1].set_xlabel(f\"spatial task distance\")\n","ax[0].set_ylabel(f\"message distance\")\n","ax[1].set_ylabel(f\"message distance\")\n","\n","for i in range(2):\n"," ax[i].set_xlim(xlims[i])\n"," ax[i].set_ylim(ylims)\n"," ax[i].tick_params(width=5, length=10)\n","\n","ax[0].set_title(\"teacher meaning\", size=28, y=1.05)\n","ax[1].set_title(\"spatial meaning\", size=28, y=1.05)\n","\n","fig.tight_layout() #only titles for columns and rows\n","\n","plt.subplots_adjust(wspace=0.5, hspace=0.6)\n","\n","plt.show()"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"vp6FYtsxLdDt"},"outputs":[],"source":["#@title 26. EXECUTION - Entropy analysis: binning approach\n","'''\n","Calculate entropy of message and other distributions\n","Always project to 2D-PCA space first (this loses some information, but makes figures comparable and also\n","reduces noise as we have a relatively small number of total samples)\n","'''\n","\n","def get_2D_PCA_entropy(input_list, pca_type, hist_method, n_Hbins):\n"," #do PCA on all inputs, extract first two coordinates\n"," input_pca=pca_type.fit_transform(input_list)\n"," xvals, yvals=input_pca[:,0], input_pca[:,1]\n"," #two different methods for getting the histogram: A or B\n"," if hist_method==\"A\":\n"," hist, xedges, yedges=np.histogram2d(xvals, yvals, bins=n_Hbins)\n"," elif hist_method==\"B\":\n"," delta = max(max(xvals)-min(xvals), max(yvals)-min(yvals))/2\n"," xmid, ymid = (max(xvals)+min(xvals))/2, (max(yvals)+min(yvals))/2\n"," hist, xedges, yedges=np.histogram2d(xvals, yvals, bins=n_Hbins, range=[[xmid-delta, xmid+delta],[ymid-delta, ymid+delta]])\n"," #get probabilities, calculate entropy and store value\n"," hist=hist.flatten()/np.sum(hist)\n"," return sp.stats.entropy(hist)\n","\n","\n","#the number of bins per dimension (so have total of n_bins**2 bins)\n","bll=len(H_binlist)\n","\n","#for messages we average entropy over 5 languages each, but for tasks and teacher Q-matrices we do not\n","H_nofeedback, H_feedback, H_feedback_noregu=np.zeros(shape=(2,bll,5)), np.zeros(shape=(2,bll,5)), np.zeros(shape=(2,bll,5))\n","H_qstd, H_pstd, H_qstd_noregu, H_pstd_noregu=np.zeros(shape=(2,bll,5)), np.zeros(shape=(2,bll,5)), np.zeros(shape=(2,bll,5)), np.zeros(shape=(2,bll,5))\n","H_tasks, H_qteacher=np.zeros(shape=(2,bll)), np.zeros(shape=(2,bll))\n","\n","pca_mssg, pca_mat, pca_task=PCA(K), PCA(grid_dim**2*4), PCA(2)\n","\n","#other preparations\n","cmap=plt.get_cmap(\"bwr\")\n","state_tensors=get_state_tensors(1) #for autoencoder forward function\n","softy=nn.Softmax(dim=0)\n","\n","#1.Entropy of the original tasks and Q-matrices\n","task_list=np.array([[i,k] for [i,j,k,l] in list(label_dict.values())])\n","q_matrix_list=np.array([torch.flatten(Q).detach().numpy() for Q in q_matrix_dict.values()])\n","\n","for k,n_Hbins in enumerate(H_binlist):\n"," for input_list, pca_type, H_arr in zip([q_matrix_list, task_list], [pca_mat, pca_task], [H_qteacher, H_tasks]):\n"," H_arr[0][k]=get_2D_PCA_entropy(input_list, pca_type, \"A\", n_Hbins)\n"," H_arr[1][k]=get_2D_PCA_entropy(input_list, pca_type, \"B\", n_Hbins)\n","\n","\n","#2.Entropy of the messages (nofeedback, feedback, feedback without regularization) and student matrices\n","nofeedback_codes=[f\"nonlinear_nostudent_language{i}\" for i in range(5)]\n","feedback_codes=[f\"nonlinear_goallocs0_zeta5_language{i}\" for i in range(5)]\n","feedback_noregu_codes=[f\"nonlinear_noregularization_language{i}\" for i in range(5)]\n","H_language_codes=[nofeedback_codes, feedback_codes, feedback_noregu_codes]\n","Hm_arrs=[H_nofeedback, H_feedback, H_feedback_noregu]\n","Hq_arrs=[None, H_qstd, H_qstd_noregu]\n","Hp_arrs=[None, H_pstd, H_pstd_noregu]\n","\n","for Hi_language_codes, Hm_arr, Hq_arr, Hp_arr in zip(H_language_codes, Hm_arrs, Hq_arrs, Hp_arrs):\n"," for l,lcode in enumerate(Hi_language_codes):\n"," message_list, autoencoder_H=get_message_list(ConvAutoEncoder, lcode, q_matrix_dict)\n"," if Hq_arr is not None:\n"," q_matrixstd_list=torch.stack([torch.flatten(autoencoder_H.student(torch.tensor([m]),state_tensors)[0])for m in message_list]).detach().numpy()\n"," p_matrixstd_list=torch.stack([torch.flatten(softy(autoencoder_H.student(torch.tensor([m]),state_tensors)[0])) for m in message_list], dim=0).detach().numpy()\n"," for k,n_Hbins in enumerate(H_binlist):\n"," Hm_arr[0][k][l]=get_2D_PCA_entropy(message_list, pca_mssg, \"A\", n_Hbins)\n"," Hm_arr[1][k][l]=get_2D_PCA_entropy(message_list, pca_mssg, \"B\", n_Hbins)\n"," if Hq_arr is not None:\n"," for input_list, H_arr in zip([q_matrixstd_list, p_matrixstd_list],[Hq_arr, Hp_arr]):\n"," H_arr[0][k][l]=get_2D_PCA_entropy(input_list, pca_mat, \"A\", n_Hbins)\n"," H_arr[1][k][l]=get_2D_PCA_entropy(input_list, pca_mat, \"B\", n_Hbins)\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"qNSMsQICarW5"},"outputs":[],"source":["#@title 27. PLOTS - Entropy binning approach (Fig. 3e)\n","'''\n","The plots of the entropy analysis\n","'''\n","\n","cmap=plt.get_cmap(\"bwr\")\n","\n","#two plots for the two different histogram methods\n","for hplot, method in zip([1], [\"equal extent\"]):\n"," #I figured out that method 2 is the better one (always lower fit errors, more reasonable...)\n"," if hplot==0:\n"," continue\n"," #calculate mean and standard deviations over the 5 languages\n"," #change hear what should be plotted (also have to change colors and labels below of course)\n"," H_arrs=[H_feedback, H_nofeedback, H_pstd, H_feedback_noregu, H_pstd_noregu]\n"," H_means, H_stds=[], []\n"," for H_arr in H_arrs:\n"," H_stds+=[np.std(H_arr[hplot], axis=1)]\n"," H_means+=[np.mean(H_arr[hplot], axis=1)]\n","\n"," #plot theoretical maximum entropy as dashed line\n"," Hmax=sp.stats.entropy([1/len(q_matrix_dict)]*len(q_matrix_dict))\n"," plt.plot([0,64], [Hmax, Hmax], ls=\"--\", color=\"black\", label=\"maximum possible entropy\", lw=3)\n"," #scatter plots of entropies and corresponding hyperbolic tangent curve fits\n"," H_fitlist=[H_qteacher[hplot]]+H_means\n"," Hstd_fitlist=[None]+H_stds\n"," colors=[\"black\", cmap(0.12), cmap(0.99), cmap(0.35), cmap(0.12), cmap(0.35)]\n"," labels=[\"teacher Q-matrices\", \"messages w/ feedback\", \"messages w/o feedback\",\"student action prob. matrices\", \"messages w/ feedback (no rec.)\", \"student action prob. matrices (no rec.)\"]\n"," markers=[\"o\", \"D\", \"s\", \"D\", \"x\", \"x\"]\n"," err=0\n"," for H, H_std, color, label, marker in zip(H_fitlist, Hstd_fitlist, colors, labels, markers):\n"," #popt, pcov, infodict, _, _=curve_fit(lambda x, a, c: tanhfunc(x, a, c, Hmax), H_binlist, H, p0=[0.5,0.5], sigma=H_std, absolute_sigma=True, full_output=True)\n"," #err+=np.sum(np.absolute(infodict[\"fvec\"]))\n"," #aparam, cparam=round(popt[0],3), round(popt[1],3)\n"," #print(aparam, cparam)\n"," #plt.plot(H_binlist, tanhfunc(H_binlist, popt[0], popt[1], Hmax), color=color, lw=3)\n"," plt.scatter(H_binlist, H, marker=marker, s=75, color=color, label=label, lw=3)\n"," #plt.errorbar(H_binlist, H, yerr=H_std, fmt=marker, ms=8, color=color, label=label, lw=3)\n","\n"," #print(f\"error is {err}\")\n","\n"," plt.legend(loc=4, fontsize=11)\n","\n"," plt.xlim(0,64)\n"," plt.ylim(2,5.5)\n"," plt.xlabel(r\"bins per PC-axis ($\\it{n}$)\")\n"," plt.ylabel(\"Shannon entropy\")\n"," plt.title(\"entropy at different comm. stages\", size=28, y=1.05)\n","\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," plt.rc('axes', labelsize=28)\n"," plt.rc('xtick', labelsize=28)\n"," plt.rc('ytick', labelsize=28)\n","\n"," plt.tick_params(width=5, length=10)\n","\n"," plt.show()\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"fv0K02hs4P8S"},"outputs":[],"source":["#@title 28. EXECUTION - Hyperparameter study (message length, student loss, number of epochs): extract rates\n","\n","'''\n","Hyperparameter analysis - we kept everything constant and varied only one of the three parameters mentioned in the title\n","'''\n","\n","#label dictionaries for all trained and unknown mazes respectively\n","label_dict_train=read_dict_from_pkl(file_loc+\"teacher/label dictionaries/\"+f\"q_matrices_labelstraining_4x4.pkl\")\n","label_dict_test=read_dict_from_pkl(file_loc+\"teacher/label dictionaries/\"+f\"q_matrices_labelstest_4x4.pkl\")\n","\n","for folder_tag, xaxis in zip([\"epochs\", \"zeta\", \"mlength_K\"], [[100,200,400,600,800,1000], [1,2,5,10,20], [1,2,3,4,5,6,7,8]]):\n","\n"," folders_known=[f\"nonlinear_{folder_tag}{i}_known\" for i in xaxis]\n"," folders_unknown=[f\"nonlinear_{folder_tag}{i}_unknown\" for i in xaxis]\n","\n"," language_nr_hyp=5\n"," param=2 #the stepfactor\n"," #potentially chuck out languages if the performances are bad\n"," for plot, ldict, known_addon, folders in zip([0,1], [label_dict_train, label_dict_test], [\"known\",\"unknown\"], [folders_known, folders_unknown]): #iterate over performance on known and unknown tasks\n","\n"," for j,folder in enumerate(folders):\n"," rates=[[],[],[],[]]\n","\n"," for language in range(language_nr_hyp):\n"," single_rates=[[],[],[],[]]\n"," #load the solving rates\n"," for k,student in enumerate([\"info\",\"misinfo\"]):\n"," single_rates[k]=np.loadtxt(file_loc+\"student/\"+f\"{folder}/{folder}_language{language}/\"+f\"solving_rates_{student}_no_learning_stepfactor2.txt\")\n"," single_rates[2]=np.loadtxt(file_loc+\"student/\"+f\"{folder}/{folder}_language{language}/\"+f\"solving_rates_rdwalkersmart_stepfactor2.txt\")\n"," single_rates[3]=np.loadtxt(file_loc+\"student/\"+f\"{folder}/{folder}_language{language}/\"+f\"solving_rates_rdwalker_stepfactor2.txt\")\n","\n"," for index in range(len(single_rates)):\n"," rates[index]+=[np.mean(single_rates[index])]\n","\n"," np.savetxt(file_loc+\"hyperparameters/\"+f\"rates {folder_tag} {j} {known_addon}\", np.array(rates))\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"dmUaL2qYEAaf"},"outputs":[],"source":["#@title 29. PLOTS - Hyperparameter study (Fig. S8)\n","'''\n","Hyperparameter analysis - we kept everything constant and varied only one of the three parameters mentioned in the title\n","'''\n","from matplotlib.patches import Rectangle\n","\n","#plot settings\n","plt.rc('axes', labelsize=28)\n","plt.rc('xtick', labelsize=28)\n","plt.rc('ytick', labelsize=28)\n","plt.rc('axes', titlesize=28)\n","cmap=plt.get_cmap(\"bwr\")\n","\n","folder_tags=[\"epochs\", \"zeta\", \"mlength_K\"]\n","rect_xs=[950,4,4.6]\n","rect_widths=[100,2,0.8]\n","bboxes=[[0.6,0.84], [1, 0.725], [0.44, 0.725]]\n","xlabels=[\"training epochs \"+r\"$N_{\\mathrm{epochs}}$\", \"loss weighting \"+r\"$\\zeta$\", \"message length \"+r\"$K$\"]\n","\n","x_axes=[[100,200,400,600,800,1000], [1,2,5,10,20], [1,2,3,4,5,6,7,8]]\n","for index, (folder_tag, xaxis, rect_x, rect_width, bbox, xlabel) in enumerate(zip(folder_tags, x_axes, rect_xs, rect_widths, bboxes, xlabels)):\n","\n"," #plot\n"," fig, ax=plt.subplots(1,2, figsize=(17,5))\n"," ax[0].grid(visible=True, axis=\"y\")\n"," ax[1].grid(visible=True, axis=\"y\")\n"," plt.rcParams['svg.fonttype']='none' #\"to make later editing of figures easier\" (Carlos)\n"," labels=[\"informed student\",\"misinformed student\",\"smart random walker\",\"random walker\"]\n"," colors=[cmap(0.15),cmap(0.25),cmap(0.99),cmap(0.75)]\n","\n"," language_nr_hyp=5\n"," param=2 #the stepfactor\n"," #potentially chuck out languages if the performances are bad\n"," for plot, known_addon in zip([0,1], [\"known\",\"unknown\"]): #iterate over performance on known and unknown tasks\n"," avg_rates=[[],[],[],[]]\n"," std_rates=[[],[],None, None]\n","\n"," for j,_ in enumerate(xaxis):\n","\n"," rates=np.loadtxt(file_loc+\"hyperparameters/\"+f\"rates {folder_tag} {j} {known_addon}\")\n"," #print(f\"info: {np.round(100*np.array(rates[0]),1)} - average {round(100*np.mean(rates[0]),1)}\")\n"," #print(f\"misinfo: {np.round(100*np.array(rates[1]),1)} - average {round(100*np.mean(rates[1]),1)}\")\n"," #print(f\"smart rd walker: {round(100*rates[2][0],1)}\")\n"," #print(f\"rd walker: {round(100*rates[3][0],1)}\")\n","\n"," #add average solving rate and NOW STANDARD ERROR OF THE MEAN\n"," for j,rates_student in enumerate(rates):\n"," avg_rates[j]+=[100*np.mean(rates_student)]\n"," if j<2:\n"," std_rates[j]+=[100*np.std(rates_student, ddof=1)]\n"," #sem_rates[j]+=[100*np.std(rates_student, ddof=1)/mt.sqrt(language_nr_hyp)]\n","\n"," for yrates,std,label,color in zip(avg_rates, std_rates, labels, colors):\n"," ax[plot].plot(xaxis, yrates, color=color, lw=3)\n"," ax[plot].scatter(xaxis, yrates, color=color, label=label, marker=\"X\", s=150)\n"," #ax[plot].errorbar(xaxis, yrates, yerr=std, color=color, label=label, fmt=\"X\", ms=15, lw=3)\n","\n"," ax[plot].set_ylim(0,100)\n"," ax[plot].set_ylabel(\"tasks solved (%)\")\n"," ax[plot].set_xlim(0,1.1*xaxis[-1])\n","\n"," #Define the coordinates and dimensions of the rectangle\n"," rect_height = 98 # height of the rectangle (can be adjusted as needed)\n"," # Create a Rectangle patch\n"," rect1 = Rectangle((rect_x, 1), rect_width, rect_height, fill=False, edgecolor='black', label=\"value in main results\", lw=3)\n"," rect2 = Rectangle((rect_x, 1), rect_width, rect_height, fill=False, edgecolor='black', label=\"value in main results\", lw=3)\n"," # Add the Rectangle patch to the plot\n"," if plot==0:\n"," ax[plot].add_patch(rect1)\n"," elif plot==1:\n"," ax[plot].add_patch(rect2)\n","\n"," ax[plot].set_yticks([0,20,40,60,80,100])\n"," ax[plot].set_xlabel(xlabel)\n"," if index==0 and plot==0:\n"," ax[plot].legend(loc='upper right', bbox_to_anchor=bbox, fontsize=15)\n","\n"," ax[plot].tick_params(width=5, length=10)\n","\n"," ax[0].set_title(f\"trained tasks\", y=1.05)\n"," ax[1].set_title(f\"test tasks\", y=1.05)\n","\n"," plt.subplots_adjust(wspace=0.5, hspace=0.6)\n","\n"," plt.show()\n"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# p-values for Fig. S3 (columns identical to Tab. 6)\n","# Figure maze scenario p-val (Info vs misinfo) p-val (Info vs Smart Random)\n","# S3a i-vii <0.0005 <0.0005\n","# \n","# S3b ii <0.0005 <0.0005\n","# iii <0.0005 <0.0005\n","# iv n.s. n.s.\n","# v n.s. n.s.\n","# vi n.s. n.s.\n","# vii n.s. n.s.\n","# \n","# S3c i-vii <0.0005 <0.0005\n","# \n","# S3d ii <0.0005 <0.0005\n","# iii <0.0005 <0.0005\n","# iv n.s. n.s.\n","# v n.s. n.s.\n","# vi n.s. n.s.\n","# vii n.s. n.s. "]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# p-values for Fig. S4 (columns identical to Tab. 6)\n","# Figure maze scenario p-val (Info vs misinfo) p-val (Info vs Smart Random)\n","# S4a i-vii <0.0005 <0.0005\n","# \n","# S4b ii <0.0005 <0.0005\n","# iii <0.0005 <0.0005\n","# iv n.s. n.s.\n","# v <0.0005 0.001\n","# vi n.s. n.s.\n","# vii 0.004 n.s.\n","# \n","# S4c i-vii <0.0005 <0.0005\n","# \n","# S4d ii <0.0005 <0.0005\n","# iii <0.0005 <0.0005\n","# iv n.s. n.s.\n","# v <0.0005 0.005\n","# vi n.s. n.s.\n","# vii <0.0005 0.013"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["# p-values for Fig. S13 (columns identical to Tab. 6)\n","# Figure maze scenario p-val (Info vs misinfo) p-val (Info vs Smart Random)\n","# S13b i 0.006 n.s.\n","# ii 0.002 0.005\n","# iii n.s. n.s.\n","# iv n.s. <0.0005\n","# v n.s. n.s.\n","# vi 0.019 <0.0005\n","# vii n.s. 0.009\n","# \n","# S13b ii 0.004 n.s.\n","# iii n.s. n.s.\n","# iv n.s. n.s.\n","# v n.s. n.s.\n","# vi n.s. n.s.\n","# vii n.s. n.s.\n","# \n","# S13c i 0.003 n.s.\n","# ii 0.001 n.s.\n","# iii n.s. n.s.\n","# iv n.s. 0.004\n","# v n.s. n.s.\n","# vi 0.021 0.007\n","# vii n.s. n.s.\n","# \n","# S13d ii 0.004 n.s.\n","# iii n.s. n.s.\n","# iv n.s. n.s.\n","# v n.s. n.s.\n","# vi n.s. n.s.\n","# vii n.s. n.s."]}],"metadata":{"colab":{"authorship_tag":"ABX9TyM/l4PXzJa+x+VlROe33oXZ","provenance":[{"file_id":"1bIVrovkoAPUKY4o-avGCj6NEaI0-Lzwl","timestamp":1702628566706},{"file_id":"1-v_AFJCShlNsusYfTZFzhPX-JVV2TKsX","timestamp":1661956785099}]},"gpuClass":"standard","kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0}