def plot_dist_to_neighbours(test_env_name, sim_ind,state_reps, geodesic_dist=True, single_pos=True):
# make new env to run test in
env = gym.make(test_env_name)
plt.close()
# make graph of env states
G = make_env_graph(env)
gd= compute_graph_distance_matrix(G,env)
dist_in_state_space = np.delete(gd[sim_ind],sim_ind) #distance from sim ind to all other states
cmap = linc_coolwarm_r #cm.get_cmap('coolwarm')
try:
reps_as_matrix = np.zeros((400,state_reps[0].shape[1]))
except:
reps_as_matrix = np.zeros((400,state_reps[0].shape[0]))
reps_as_matrix[:] = np.nan
for ind, (k,v) in enumerate(state_reps.items()):
reps_as_matrix[k] = v
if k in env.obstacle:
reps_as_matrix[k,:] = np.nan
RS = squareform(pdist(reps_as_matrix,metric='chebyshev'))
for state2d in env.obstacle:
RS[state2d,:] = np.nan
RS[:,state2d] = np.nan
dist_in_rep_space = np.delete(RS[sim_ind],sim_ind)
print(type(dist_in_rep_space))
rgba = [cmap(x) for x in dist_in_rep_space/np.nanmax(dist_in_rep_space)]
if geodesic_dist:
plt.scatter(dist_in_state_space,dist_in_rep_space,color=LINCLAB_COLS['red'],alpha=0.4,linewidths=0.5 )
plt.xlabel("D(s,s')")
plt.ylabel("D(R(s), R(s'))")
plt.ylim([0.,1.1])
plt.savefig(f'../figures/CH1/latent_distance{sim_ind}.svg')
else:
if single_pos:
a = plt.imshow(RS[sim_ind].reshape(env.shape)/np.nanmax(RS[sim_ind]), cmap=cmap, vmin=0, vmax=1)
r,c = env.oneD2twoD(sim_ind)
plt.gca().add_patch(plt.Rectangle(np.add((c,r),(-0.5,-0.5)), .99, .99, edgecolor='k', fill=False, alpha=1))
plt.colorbar(a)
plt.gca().get_xaxis().set_visible(False)
plt.gca().get_yaxis().set_visible(False)
plt.savefig(f'../figures/CH1/representation_similarity/latent_sim{sim_ind}.svg')
else:
sliced = RS.copy()
for state1d in reversed(env.obstacle):
sliced = np.delete(sliced,state1d,0)
sliced = np.delete(sliced,state1d,1)
a = plt.imshow(sliced/np.nanmax(sliced), vmin=0, vmax=1, cmap=linc_coolwarm)
plt.colorbar(a)
#plt.show()
plt.close()
def get_graph_dist_from_state(envs_to_plot, sim_ind):
graph_distances = []
for i, test_env_name in enumerate(envs_to_plot):
env = gym.make(test_env_name)
plt.close()
G = make_env_graph(env)
gd = compute_graph_distance_matrix(G, env)
dist_in_state_space = gd[sim_ind]
graph_distances.append(dist_in_state_space)
return graph_distances
def attempt_opt_pol(env):
rwd_loc = env.twoD2oneD(list(env.rewards.keys())[0])
G = make_env_graph(env)
sp = nx.shortest_path(G)
gd = compute_graph_distance_matrix(G, env)
num_steps_to_rwd = np.zeros((20, 20))
num_steps_to_rwd[:] = np.nan
for (i, j) in env.useable:
# i = row, j = col
oneD = env.twoD2oneD((i, j))
num_steps_to_rwd[i, j] = gd[oneD, rwd_loc]
opt_pol_matrix = np.zeros((20, 20, 4))
opt_pol_matrix[:] = np.nan
for (r, c) in env.useable:
index = env.twoD2oneD((r, c))
steps_from_index = num_steps_to_rwd[r, c]
if r + 1 < 20:
state_down = (r + 1, c)
steps_down = num_steps_to_rwd[state_down]
if steps_down < steps_from_index:
opt_pol_matrix[r, c, 0] = 1
if r - 1 >= 0:
state_up = (r - 1, c)
index_up = env.twoD2oneD(state_up)
steps_up = num_steps_to_rwd[state_up]
if steps_up < steps_from_index:
opt_pol_matrix[r, c, 1] = 1
if c + 1 < 20:
state_right = (r, c + 1)
steps_right = num_steps_to_rwd[state_right]
if steps_right < steps_from_index:
opt_pol_matrix[r, c, 2] = 1
if (c - 1) >= 0:
state_left = (r, c - 1)
steps_left = num_steps_to_rwd[state_left]
if steps_left < steps_from_index:
opt_pol_matrix[r, c, 3] = 1
opt_pol_matrix[r,
c, :] = np.nan_to_num(softmax(opt_pol_matrix[r, c, :]))
return opt_pol_matrix
#opt_pol = attempt_opt_pol(env)
#op = opt_pol_map(gym.make('gridworld:gridworld-v1'))
#op[9:,:] = [0., 0.,0,0]
#plot_pref_pol(env, opt_pol)
def plot_grid_of_shit(test_env_name,state_reps, sim_ind, metric):
# make new env to run test in
env = gym.make(test_env_name)
plt.close()
# make graph of env states
G = make_env_graph(env)
gd= compute_graph_distance_matrix(G,env)
dist_in_state_space = np.delete(gd[sim_ind],sim_ind) #distance from sim ind to all other states
cmap = linc_coolwarm_r #cmx.get_cmap('Spectral_r') #
try:
reps_as_matrix = np.zeros((400,state_reps[0].shape[1]))
except:
reps_as_matrix = np.zeros((400,state_reps[0].shape[0]))
reps_as_matrix[:] = np.nan
for ind, (k,v) in enumerate(state_reps.items()):
reps_as_matrix[k] = v
if k in env.obstacle:
reps_as_matrix[k,:] = np.nan
RS = squareform(pdist(reps_as_matrix,metric=metric))
for state2d in env.obstacle:
RS[state2d,:] = np.nan
RS[:,state2d] = np.nan
dist_in_rep_space = np.delete(RS[sim_ind],sim_ind)
a = plt.imshow(RS[sim_ind].reshape(env.shape)/np.nanmax(RS[sim_ind]), cmap=cmap)
r,c = env.oneD2twoD(sim_ind)
plt.gca().add_patch(plt.Rectangle(np.add((c,r),(-0.5,-0.5)), .99, .99, edgecolor='k', fill=False, alpha=1))
plt.colorbar(a)
plt.gca().get_xaxis().set_visible(False)
plt.gca().get_yaxis().set_visible(False)
plt.savefig(f'../figures/CH1/representation_similarity/distance_metrics/{test_env_name[-2:]}{rep_name}{metric}.svg')
plt.show()
for id_num in id_list:
with open(f'../../Data/results/{id_num}_data.p', 'rb') as f:
data = pickle.load(f)
state_occ_vec = data['occupancy']
all_occ += state_occ_vec
all_visits = np.sum(state_occ_vec)
all_v.append(all_visits)
print(all_occ)
state_occ_map[:] = np.nan
for r,c in env.useable:
ind = env.twoD2oneD((r,c))
state_occ_map[r,c] = all_occ[ind]
G = make_env_graph(env)
gd = compute_graph_distance_matrix(G, env)
dist_from_reward = gd[:,105]
#a = plt.imshow(state_occ_map/(all_visits/len(env.useable)),vmax=4)
#plt.colorbar(a)
#plt.show()
print(np.nanmax(dist_from_reward))
janky_histo = np.zeros(int(np.nanmax(dist_from_reward))+1)
num_times_for_dist = np.zeros(int(np.nanmax(dist_from_reward))+1)
for ind in range(len(state_occ_vec)):
print(dist_from_reward[ind])
if np.isnan(dist_from_reward[ind]):
pass
def plot_world(world, **kwargs):
scale = kwargs.get('scale', 0.35)
title = kwargs.get('title', 'Grid World')
ax_labels = kwargs.get('ax_labels', False)
state_labels = kwargs.get('states', False)
invert_ = kwargs.get('invert', False)
if invert_:
cmap = 'bone'
else:
cmap = 'bone_r'
r, c = world.shape
G = make_env_graph(env)
sp = nx.shortest_path(G)
gd = compute_graph_distance_matrix(G, world)
fig = plt.figure(figsize=(c * scale, r * scale))
ax = fig.add_subplot(1, 1, 1)
gridMat = np.zeros(world.shape)
for i, j in world.obstacle2D:
gridMat[i, j] = 1.0
for i, j in world.terminal2D:
gridMat[i, j] = 0.2
ax.pcolor(world.grid,
edgecolors='k',
linewidths=0.75,
cmap=cmap,
vmin=0,
vmax=1)
U = np.zeros((r, c))
V = np.zeros((r, c))
U[:] = np.nan
V[:] = np.nan
if len(world.action_list) > 4:
if world.jump is not None:
for (a, b) in world.jump.keys():
(a2, b2) = world.jump[(a, b)]
U[a, b] = (b2 - b)
V[a, b] = (a - a2)
C, R = np.meshgrid(np.arange(0, c) + 0.5, np.arange(0, r) + 0.5)
ax.quiver(C, R, U, V, scale=1, units='xy')
for rwd_loc in world.rewards.keys():
rwd_r, rwd_c = rwd_loc
if world.rewards[rwd_loc] < 0:
colorcode = 'red'
else:
colorcode = 'darkgreen'
ax.add_patch(
plt.Rectangle((rwd_c, rwd_r),
width=1,
height=1,
linewidth=2,
facecolor=colorcode,
alpha=0.5))
if state_labels:
for (i, j) in world.useable:
# i = row, j = col
oneD = world.twoD2oneD((i, j))
#ax.text(j+0.5,i+0.7, s=f'{oneD}', ha='center')
ax.text(j + 0.5, i + 0.7, s=f'{gd[oneD,105]}', ha='center')
#ax.set_xticks([np.arange(c) + 0.5, np.arange(c)])
#ax.set_yticks([np.arange(r) + 0.5, np.arange(r)])
ax.invert_yaxis()
ax.set_aspect('equal')
if not ax_labels:
ax.get_xaxis().set_ticks([])
ax.get_yaxis().set_ticks([])
ax.set_title(title)
return fig, ax
def plot_dist_to_neighbours():
for test_env_name in [envs_to_plot[1]]:
sim_ind = 169
envno = envs_to_plot.index(test_env_name)
# make new env to run test in
env = gym.make(test_env_name)
plt.close()
# make graph of env states
G = make_env_graph(env)
gd = compute_graph_distance_matrix(G, env)
dist_in_state_space = np.delete(
gd[sim_ind], sim_ind) #distance from sim ind to all other states
rep_types = {
'random': random,
'onehot': onehot,
'place_cell': place_cell,
'sr': sr
} #, 'latents':latents}
fig, ax = plt.subplots(1,
len(list(rep_types.items())) + 1,
figsize=(14, 2))
if test_env_name[-2:] == '51':
rwd_colrow = (16, 9)
else:
rwd_colrow = (14, 14)
rect = plt.Rectangle(rwd_colrow, 1, 1, color='g', alpha=0.3)
agt_colrow = (env.oneD2twoD(sim_ind)[1] + 0.5,
env.oneD2twoD(sim_ind)[0] + 0.5)
circ = plt.Circle(agt_colrow, radius=0.3, color='blue')
ax[0].pcolor(grids[envno],
cmap='bone_r',
edgecolors='k',
linewidths=0.1)
ax[0].axis(xmin=0, xmax=20, ymin=0, ymax=20)
ax[0].set_aspect('equal')
ax[0].add_patch(rect)
ax[0].add_patch(circ)
ax[0].get_xaxis().set_visible(False)
ax[0].get_yaxis().set_visible(False)
ax[0].invert_yaxis()
cmap = linc_coolwarm_r #cm.get_cmap('coolwarm')
for j, rep_type in enumerate(rep_types.keys()):
relative_path_to_data = '../../Data/' # from within Tests/CH1
if rep_type == 'latents':
conv_ids = {
'gridworld:gridworld-v1':
'c34544ac-45ed-492c-b2eb-4431b403a3a8',
'gridworld:gridworld-v3':
'32301262-cd74-4116-b776-57354831c484',
'gridworld:gridworld-v4':
'b50926a2-0186-4bb9-81ec-77063cac6861',
'gridworld:gridworld-v5':
'15b5e27b-444f-4fc8-bf25-5b7807df4c7f'
}
run_id = conv_ids[test_env_name[:-1]]
agent_path = relative_path_to_data + f'agents/{run_id}.pt'
state_reps, representation_name, input_dims, _ = latents(
env, agent_path)
else:
state_reps, representation_name, input_dims, _ = rep_types[
rep_type](env)
reps_as_matrix = np.zeros((400, 400))
reps_as_matrix[:] = np.nan
for ind, (k, v) in enumerate(state_reps.items()):
reps_as_matrix[k] = v
if k in env.obstacle:
reps_as_matrix[k, :] = np.nan
RS = squareform(pdist(reps_as_matrix, metric='chebyshev'))
for state2d in env.obstacle:
RS[state2d, :] = np.nan
RS[:, state2d] = np.nan
dist_in_rep_space = np.delete(RS[sim_ind], sim_ind)
print(type(dist_in_rep_space))
rgba = [
cmap(x)
for x in dist_in_rep_space / np.nanmax(dist_in_rep_space)
]
ax[j + 1].scatter(dist_in_state_space,
dist_in_rep_space,
color='#b40426',
alpha=0.2,
linewidths=0.5)
ax[j + 1].set_xlabel("D(s,s')")
ax[j + 1].set_ylim([0.4, 1.1])
if j != 0:
ax[j + 1].set_yticklabels([])
else:
ax[j + 1].set_ylabel("D(R(s), R(s'))")
#a = ax[j+2].imshow((dist_in_state_space/np.nanmax(dist_in_state_space)).reshape(env.shape), cmap=cmap)
#plt.colorbar(a, ax=ax[j+2])
plt.savefig('../figures/CH2/distance.svg')
plt.show()