def ssp_enc_func(coords, dim=256, seed=13): rng = np.random.RandomState(seed) x_axis_sp = make_good_unitary(dim, rng=rng) y_axis_sp = make_good_unitary(dim, rng=rng) return encode_point( x=coords[0], y=coords[1], x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp, ).v
def test_generate_maze_sp(size=10, limit_low=-5, limit_high=5, res=64, dim=512, seed=13): from spatial_semantic_pointers.utils import make_good_unitary, get_heatmap_vectors from spatial_semantic_pointers.plots import plot_heatmap rng = np.random.RandomState(seed=seed) x_axis_sp = make_good_unitary(dim=dim, rng=rng) y_axis_sp = make_good_unitary(dim=dim, rng=rng) xs = np.linspace(limit_low, limit_high, res) ys = np.linspace(limit_low, limit_high, res) sp, maze, fine_maze = generate_maze_sp(size, xs, ys, x_axis_sp, y_axis_sp, normalize=True, obstacle_ratio=.2, map_style='blocks') fig, ax = plt.subplots(1, 4) ax[0].imshow(maze) ax[1].imshow(fine_maze) heatmap_vectors = get_heatmap_vectors(xs, ys, x_axis_sp, y_axis_sp) plot_heatmap(sp.v, heatmap_vectors, ax[2], xs, ys, name='', vmin=-1, vmax=1, cmap='plasma', invert=True) plot_heatmap(sp.v, heatmap_vectors, ax[3], xs, ys, name='', vmin=None, vmax=None, cmap='plasma', invert=True) plt.show()
def debug(): seed = 14 dim = 256 # dim = 1024 res = 256 #32 res = 32 rng = np.random.RandomState(seed=seed) X, Y = get_fixed_dim_sub_toriod_axes( dim=dim, n_proj=3, scale_ratio=0, scale_start_index=0, rng=rng, eps=0.001, ) X, Y = get_axes_from_network(fname='learned_ssp_models/model_1layer_1024hs_seed1.pt', dim=dim) # X, Y = get_axes(dim=dim, n=3, seed=13, period=0, optimal_phi=False) X = make_good_unitary(dim, rng=rng) Y = make_good_unitary(dim, rng=rng) # X = make_fixed_dim_periodic_axis( # dim=dim, period=6, phase=0, frequency=1, eps=1e-3, rng=rng, flip=False, random_phases=False, # ) # Y = make_fixed_dim_periodic_axis( # dim=dim, period=6, phase=0, frequency=1, eps=1e-3, rng=rng, flip=False, random_phases=False, # ) md, phi_pos = phi_mag_and_dir(X.v, Y.v) # md, phi_pos = phi_mag_and_dir(X, Y) fig, ax = plt.subplots(1, 5, figsize=(10, 3)) magnitude_histogram(md, ax[0]) direction_histogram(md, ax[1]) mag_dir_histogram(md, ax[2], res=res) mag_dir_histogram(phi_pos, ax[3], res=res) gauss_image(phi_pos, ax[4], sigma=.15) print("Mean Location: ({},{})".format( np.round(np.mean(phi_pos[:, 0]), 2), np.round(np.mean(phi_pos[:, 1]), 2) )) plt.show()
def ssp_encoding_func(seed=13, dim=512, ssp_scaling=1): rng = np.random.RandomState(seed=seed) x_axis_sp = make_good_unitary(dim=dim, rng=rng) y_axis_sp = make_good_unitary(dim=dim, rng=rng) def encoding_func(positions): return encode_point( x=positions[0]*ssp_scaling, y=positions[1]*ssp_scaling, x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp ).v return encoding_func
def encode_dataset(data, dim=256, seed=13, scale=1.0): """ :param data: the data to be encoded :param dim: dimensionality of the SSP :param seed: seed for the single axis vector :param scale: scaling of the data for the encoding :return: """ rng = np.random.RandomState(seed=seed) # TODO: have option to normalize everything first, for consistent relative scale axis_vec = make_good_unitary(dim, rng=rng) n_samples = data.shape[0] n_features = data.shape[1] n_out_features = n_features * dim data_out = np.zeros((n_samples, n_out_features)) for s in range(n_samples): for f in range(n_features): data_out[s, f * dim:(f + 1) * dim] = power(axis_vec, data[s, f] * scale).v return data_out
def get_nd_simplex_encoding_func(n, dim, seed=13): transform_axes = get_simplex_coordinates(n) rng = np.random.RandomState(seed=seed) axis_vectors = [] for i in range(n + 1): axis_vectors.append(make_good_unitary(dim, rng=rng)) def encoding_func(features): """ Take in 'n' features as a numpy array, and output a 'dim' dimensional SSP """ # TODO: any scaling required? # TODO: make sure matrix multiply order is correct tranformed_features = transform_axes @ features vec = power(axis_vectors[0], tranformed_features[0]) for i in range(1, n + 1): vec *= power(axis_vectors[i], tranformed_features[i]) return vec.v return encoding_func
def random_unitary(n_samples=1000, dim=3, version=1, eps=0.001): points = np.zeros((n_samples, dim)) good = np.zeros((n_samples, )) for i in range(n_samples): if version == 1: sp = nengo_spa.SemanticPointer(data=np.random.randn(dim)) sp = sp.normalized() sp = sp.unitary() elif version == 0: sp = spa.SemanticPointer(dim) sp.make_unitary() elif version == 2: sp = make_good_unitary(dim=dim) else: raise NotImplementedError points[i, :] = sp.v pf = np.fft.fft(points[i, :]) if dim % 2 == 0: if np.abs(pf[0] - 1) < eps and np.abs(pf[dim // 2] - 1) < eps: good[i] = 1 else: if np.abs(pf[0] - 1) < eps: good[i] = 1 return points, good
def unitary_points(dim, n_samples, good_unitary=False): points = np.zeros((n_samples, dim)) for i in range(n_samples): if good_unitary: points[i, :] = make_good_unitary(dim=dim).v else: sp = spa.SemanticPointer(dim) sp.make_unitary() points[i, :] = sp.v return points
def get_ssp_encoding_func(dim, scale, seed, **_): rng = np.random.RandomState(seed=seed) axis_vec = make_good_unitary(dim, rng=rng) def encoding_func(feature): return power(axis_vec, feature * scale).v return encoding_func
def unitary_determinant_test(dim=2, n_samples=1000): for i in range(n_samples): if True: vec = make_good_unitary(dim=dim).v print(vec) else: sp = spa.SemanticPointer(dim) sp.make_unitary() vec = sp.v if not np.allclose(np.dot(circulant(vec)[0], circulant(vec)[1]), 0): print(np.dot(circulant(vec)[0], circulant(vec)[1])) print(np.linalg.det(circulant(vec))) # assert np.allclose(np.dot(circulant(vec)[0], circulant(vec)[1]), 0) assert np.allclose(np.abs(np.linalg.det(circulant(vec))), 1) return True
def get_heatmap_vectors_n(xs, ys, n, seed=13, dim=512): """ Precompute spatial semantic pointers for every location in the linspace Used to quickly compute heat maps by a simple vectorized dot product (matrix multiplication) """ rng = np.random.RandomState(seed=seed) axis_sps = [] for i in range(n): axis_sps.append(make_good_unitary(dim, rng=rng)) vectors = np.zeros((len(xs), len(ys), dim)) N = len(axis_sps) # points_nd = np.zeros((N + 1, N)) # points_nd[:N, :] = np.eye(N) # # points in 2D that will correspond to each axis, plus one at zero # points_2d = np.zeros((N + 1, 2)) points_nd = np.eye(N) * np.sqrt(N) # points in 2D that will correspond to each axis points_2d = np.zeros((N, 2)) thetas = np.linspace(0, 2 * np.pi, N + 1)[:-1] # TODO: will want a scaling here, or along the high dim axes for i, theta in enumerate(thetas): points_2d[i, 0] = np.cos(theta) points_2d[i, 1] = np.sin(theta) transform_mat = np.linalg.lstsq(points_2d, points_nd) # apply scaling to the axes based on the singular values. Both should be the same x_axis = transform_mat[0][0, :] / transform_mat[3][0] y_axis = transform_mat[0][1, :] / transform_mat[3][1] for i, x in enumerate(xs): for j, y in enumerate(ys): # Note: needed to divide by sqrt(2) in the scaling here to get it to match with the 2D and regular hex 3D method. p = encode_point_n(x=x * transform_mat[3][0] / np.sqrt(2), y=y * transform_mat[3][0] / np.sqrt(2), axis_sps=axis_sps, x_axis=x_axis, y_axis=y_axis) vectors[i, j, :] = p.v # also return the axis_sps so individual points can be generated return vectors, axis_sps
def random_unitary(n_samples=1000, dim=3, version=2): points = np.zeros((n_samples, dim)) for i in range(n_samples): if version == 1: sp = nengo_spa.SemanticPointer(data=np.random.randn(dim)) sp = sp.normalized() sp = sp.unitary() elif version == 0: sp = spa.SemanticPointer(dim) sp.make_unitary() elif version == 2: sp = make_good_unitary(dim=dim) else: raise NotImplementedError points[i, :] = sp.v return points
def get_nd_encoding_func(n, dim, seed=13): rng = np.random.RandomState(seed=seed) axis_vectors = [] for i in range(n): axis_vectors.append(make_good_unitary(dim, rng=rng)) def encoding_func(features): """ Take in 'n' features as a numpy array, and output a 'dim' dimensional SSP """ vec = power(axis_vectors[0], features[0]) for i in range(1, n): vec *= power(axis_vectors[i], features[i]) return vec.v return encoding_func
def get_axes_and_scale(dim=256, n=3, seed=13, apply_scaling=True): """ Get X and Y axis vectors based on an n dimensional projection. Also return the correct scaling so the size of the bump in the heatmap is consistent """ rng = np.random.RandomState(seed=seed) points_nd = np.eye(n) * np.sqrt(n) # points in 2D that will correspond to each axis, plus one at zero points_2d = np.zeros((n, 2)) thetas = np.linspace(0, 2 * np.pi, n + 1)[:-1] # TODO: will want a scaling here, or along the high dim axes for i, theta in enumerate(thetas): points_2d[i, 0] = np.cos(theta) points_2d[i, 1] = np.sin(theta) transform_mat = np.linalg.lstsq(points_2d, points_nd) x_axis = transform_mat[0][0, :] y_axis = transform_mat[0][1, :] if apply_scaling: x_axis /= transform_mat[3][0] y_axis /= transform_mat[3][1] axis_sps = [] for i in range(n): # random unitary vector axis_sps.append(make_good_unitary(dim, rng=rng)) X = power(axis_sps[0], x_axis[0]) Y = power(axis_sps[0], y_axis[0]) for i in range(1, n): X *= power(axis_sps[i], x_axis[i]) Y *= power(axis_sps[i], y_axis[i]) return X, Y, transform_mat[3][0]
version = 2 n_samples = 25000 #10000 dim = 5 #3 points = np.zeros((n_samples, dim)) for i in range(n_samples): if version == 1: sp = nengo_spa.SemanticPointer(data=np.random.randn(dim)) sp = sp.normalized() sp = sp.unitary() elif version == 0: sp = spa.SemanticPointer(dim) sp.make_unitary() elif version == 2: sp = make_good_unitary(dim=dim) points[i, :] = sp.v fig = plt.figure() ax = fig.add_subplot(111, projection='3d') def orthogonal_dir_unitary(dim=5, phi=np.pi / 2.): xf = np.zeros((dim, ), dtype='Complex64') xf[0] = 1 xf[1] = np.exp(1.j * phi) xf[2] = 1 xf[3] = 1 xf[4] = np.exp(-1.j * phi)
else: raise NotImplementedError points[i, :] = sp.v return points # unitaries = [ # unitary(np.pi/3., +1), # unitary(-np.pi/8., +1), # unitary(np.pi/3., -1), # unitary(np.pi/8., -1), # ] unitaries = [ make_good_unitary(dim=5, rng=np.random.RandomState(seed=13)).v, make_good_unitary(dim=11, rng=np.random.RandomState(seed=13)).v, make_good_unitary(dim=32, rng=np.random.RandomState(seed=14)).v, make_good_unitary(dim=128, rng=np.random.RandomState(seed=13)).v, ] titles = [ 'Dim = 5', 'Dim = 11', 'Dim = 32', 'Dim = 128', ] palette = sns.color_palette("hls", len(unitaries)) # for shading the space
def main(): parser = argparse.ArgumentParser( 'Traverse many graphs with and SSP algorithm and report metrics') parser.add_argument('--n-samples', type=int, default=5, help='Number of different graphs to test') parser.add_argument('--seed', type=int, default=13, help='Seed for training and generating axis SSPs') parser.add_argument('--dim', type=int, default=512, help='Dimensionality of the SSPs') parser.add_argument('--res', type=int, default=128, help='Resolution of the linspaces used') parser.add_argument('--normalize', type=int, default=1, choices=[0, 1], help='Whether or not to normalize SPs') parser.add_argument( '--diameter-increment', type=float, default=1.0, help='How much to expand ellipse diameter by on each step') args = parser.parse_args() # Convert to boolean args.normalize = args.normalize == 1 # Metrics # set to 1 if the found path is the shortest shortest_path = np.zeros((args.n_samples)) # set to 1 if the path found is valid (only uses connections that exist) valid_path = np.zeros((args.n_samples)) # set to 1 if any path is found in the time allotted found_path = np.zeros((args.n_samples)) np.random.seed(args.seed) xs = np.linspace(0, 10, args.res) ys = np.linspace(0, 10, args.res) x_axis_sp = make_good_unitary(args.dim) y_axis_sp = make_good_unitary(args.dim) heatmap_vectors = get_heatmap_vectors(xs, ys, x_axis_sp, y_axis_sp) # TEMP: putting this outside the loop for debugging graph_params = generate_graph(dim=args.dim, x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp) for n in range(args.n_samples): print("Sample {} of {}".format(n + 1, args.n_samples)) # graph_params = generate_graph(dim=args.dim, x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp) elliptic_expansion = EllipticExpansion( x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp, xs=xs, ys=ys, heatmap_vectors=heatmap_vectors, diameter_increment=args.diameter_increment, normalize=args.normalize, debug_mode=True, **graph_params) path = elliptic_expansion.find_path( max_steps=10, #15,#20, display=False, graph=graph_params['graph'], xs=xs, ys=ys, heatmap_vectors=heatmap_vectors) optimal_path = graph_params['graph'].search_graph( start_node=graph_params['start_landmark_id'], end_node=graph_params['end_landmark_id'], ) print("found path is: {}".format(path)) print("optimal path is: {}".format(optimal_path)) if path is not None: found_path[n] = 1 if graph_params['graph'].is_valid_path(path): valid_path[n] = 1 print("path is valid") else: print("path is invalid") if path == optimal_path: shortest_path[n] = 1 print("path is optimal") else: print("path is not optimal") print("Found path: {}".format(found_path.mean())) print("Valid path: {}".format(valid_path.mean())) print("Shortest path: {}".format(shortest_path.mean()))
xs = np.linspace(-args.limit, args.limit, args.res) ys = np.linspace(-args.limit, args.limit, args.res) # if the data already exists, just load it if os.path.exists(fname): data = np.load(fname) square_heatmaps = data['square_heatmaps'] hex_heatmaps = data['hex_heatmaps'] else: for seed in range(args.n_seeds): print("\x1b[2K\r Seed {} of {}".format(seed + 1, args.n_seeds), end="\r") rng = np.random.RandomState(seed=seed) X = make_good_unitary(args.dim, rng=rng) Y = make_good_unitary(args.dim, rng=rng) Z = make_good_unitary(args.dim, rng=rng) axes[seed, 0, :] = X.v axes[seed, 1, :] = Y.v axes[seed, 2, :] = Z.v square_heatmaps[seed, :, :] = np.tensordot(origin_vec, get_heatmap_vectors( xs=xs, ys=ys, x_axis_sp=X, y_axis_sp=Y), axes=([0], [2]))
# data = np.load(fname) # square_heatmaps = data['square_heatmaps'] # hex_heatmaps = data['hex_heatmaps'] # # avg_square_heatmap = square_heatmaps.mean(axis=0) # avg_hex_heatmap = hex_heatmaps.mean(axis=0) dim = 512 res = 256 limit = 5 xs = np.linspace(-limit, limit, res) Xh, Yh = get_axes(dim=dim, seed=13) rng = np.random.RandomState(seed=13) X = make_good_unitary(dim=dim, rng=rng) Y = make_good_unitary(dim=dim, rng=rng) Z = make_good_unitary(dim=dim, rng=rng) fig, ax = plt.subplots(1, 3, figsize=(8, 4)) sigma_normal = 0.5 sigma_hex = 0.5 sigma_hex_c = 0.5 sim_hex = np.zeros((res, )) sim_hex_c = np.zeros((res, )) # this version has axes generated together and then converted to 2D sim_normal = np.zeros((res, )) gauss_hex = gaussian_1d(0, sigma_hex, xs) gauss_hex_c = gaussian_1d(0, sigma_hex_c, xs) gauss_normal = gaussian_1d(0, sigma_normal, xs)
import nengo.spa as spa import numpy as np from spatial_semantic_pointers.utils import encode_point, make_good_unitary, get_heatmap_vectors from spatial_semantic_pointers.plots import SpatialHeatmap from spatial_semantic_pointers.networks.ssp_cleanup import SpatialCleanup seed = 13 dim = 512 limit = 5 res = 256 model = spa.SPA(seed=seed) rstate = np.random.RandomState(seed=13) x_axis_sp = make_good_unitary(dim, rng=rstate) y_axis_sp = make_good_unitary(dim, rng=rstate) xs = np.linspace(-limit, limit, res) ys = np.linspace(-limit, limit, res) n_cconv_neurons = 50 #10 #ssp_cleanup_path = '/home/bjkomer/metric-representation/metric_representation/pytorch/ssp_cleanup_cosine_15_items/May07_15-21-15/model.pt' ssp_cleanup_path = '/home/ctnuser/metric-representation/metric_representation/pytorch/ssp_cleanup_cosine_15_items/May07_15-21-15/model.pt' def angle_to_ssp(x): return encode_point(np.cos(x), np.sin(x), x_axis_sp, y_axis_sp).v
for i, x in enumerate(xs): for j, y in enumerate(ys): sim[i, j] = np.dot( point, np.fft.ifft(np.fft.fft(X)**x * np.fft.fft(Y)**y).real) return sim # X = unitary(np.pi/2., 0) # Y = unitary(0, np.pi/2.) rng = np.random.RandomState(seed=13) dim = 128 X = make_good_unitary(dim=dim, rng=rng).v Y = make_good_unitary(dim=dim, rng=rng).v # dim = 32 # X = orthogonal_unitary(dim=dim, index=dim//4, phi=np.pi/2.) # Y = orthogonal_unitary(dim=dim, index=dim//4+1, phi=np.pi/2.) pos = [.5, .25] pos = [2, 1] # representation of the point as an SSP point = np.fft.ifft(np.fft.fft(X)**pos[0] * np.fft.fft(Y)**pos[1]).real limit = 5 xs = np.linspace(-limit, limit, 64) ys = np.linspace(-limit, limit, 64)
# np.random.seed(13) # np.random.seed(17) np.random.seed(42) dim = 512 # dim = 1024 res = 128 xs = np.linspace(0, 10, res) ys = np.linspace(0, 10, res) # These will include the space that is the difference between any two nodes xs_larger = np.linspace(-10, 10, res) ys_larger = np.linspace(-10, 10, res) x_axis_sp = make_good_unitary(dim) y_axis_sp = make_good_unitary(dim) heatmap_vectors = get_heatmap_vectors(xs, ys, x_axis_sp, y_axis_sp) heatmap_vectors_larger = get_heatmap_vectors(xs_larger, ys_larger, x_axis_sp, y_axis_sp) # Map map_sp = spa.SemanticPointer(data=np.zeros((dim, ))) # version of the map with landmark IDs bound to each location landmark_map_sp = spa.SemanticPointer(data=np.zeros((dim, ))) # Connectivity # contains each connection egocentrically con_ego_sp = spa.SemanticPointer(data=np.zeros((dim, ))) # contains each connection allocentrically
if args.dim == 256: ssp_cleanup_path = '/home/ctnuser/ssp-navigation/ssp_navigation/trained_models/ssp_cleanup/hex-ssp/d256/May23_13-54-52/model.pt' elif args.dim == 512: ssp_cleanup_path = '/home/ctnuser/ssp-navigation/ssp_navigation/trained_models/ssp_cleanup/hex-ssp/d512/May23_14-48-23/model.pt' elif args.dim == 1024: ssp_cleanup_path = '/home/ctnuser/ssp-navigation/ssp_navigation/trained_models/ssp_cleanup/hex-ssp/d1024/May23_16-11-12/model.pt' else: ssp_cleanup_path = '/home/ctnuser/metric-representation/metric_representation/pytorch/ssp_cleanup_cosine_15_items/May07_15-21-15/model.pt' seed = 13 dim = 512 limit = 5 res = 256 rstate = np.random.RandomState(seed=13) X = make_good_unitary(dim, rng=rstate) Y = make_good_unitary(dim, rng=rstate) xs = np.linspace(-limit, limit, res) ys = np.linspace(-limit, limit, res) pytorch_cleanup = True if not pytorch_cleanup: # generate a dataset for the cleanup function cache_fname = 'cleanup_dataset_{}.npz'.format(dim) if os.path.exists(cache_fname): data = np.load(cache_fname) clean_vectors = data['clean_vectors'] noisy_vectors = data['noisy_vectors'] else:
if __name__ == '__main__': args = parser.parse_args() else: args = parser.parse_args([]) neurons_per_dim = 5 n_neurons = args.dim * args.neurons_per_dim xs = np.linspace(-args.limit, args.limit, args.res) ys = np.linspace(-args.limit, args.limit, args.res) # X, Y = get_axes(dim=dim, n=3, seed=13, period=0, optimal_phi=False) # temporarily using these to match the axis vectors the cleanup was trained on rstate = np.random.RandomState(seed=13) X = make_good_unitary(args.dim, rng=rstate) Y = make_good_unitary(args.dim, rng=rstate) def angle_to_ssp(x): return encode_point(np.cos(x), np.sin(x), X, Y).v heatmap_vectors = get_heatmap_vectors(xs, ys, X, Y) rng = np.random.RandomState(seed=13) preferred_locations = hilbert_2d(-args.limit, args.limit, n_neurons, rng,
neurons_per_dim = 5 n_neurons = dim * neurons_per_dim n_cconv_neurons = neurons_per_dim * 2 # preferred_locations = hilbert_2d(pc_limit_low, pc_limit_high, n_neurons, rng, p=8, N=2, normal_std=3) preferred_locations = hilbert_2d(pc_limit_low, pc_limit_high, n_neurons, rng, p=10, N=2, normal_std=3) # item_sp = nengo.spa.SemanticPointer(dim) item_sp = make_good_unitary(dim) # item_sp = encode_point(1, 1, X, Y) def input_func(t): index = int(np.floor(t / dt)) pos = positions[index % (n_samples - 1)] pos_ssp = encode_point(pos[0], pos[1], X, Y) # item = item_sp if diff_axis: bound = encode_point(pos[0], pos[1], X_new, Y_new) else: bound = item_sp * pos_ssp if index > n_samples - 1:
def get_env(seed=13, dim=512): rstate = np.random.RandomState(seed=seed) x_axis_sp = make_good_unitary(dim=dim, rng=rstate) y_axis_sp = make_good_unitary(dim=dim, rng=rstate) map_array = np.array([ [1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 0, 1, 0, 0, 1, 0, 0, 0, 1], [1, 0, 1, 1, 0, 0, 0, 0, 0, 1], [1, 0, 0, 1, 0, 0, 1, 1, 0, 1], [1, 0, 0, 0, 0, 0, 0, 1, 0, 1], [1, 0, 0, 0, 1, 1, 0, 0, 0, 1], [1, 0, 1, 0, 1, 0, 0, 0, 0, 1], [1, 1, 1, 0, 0, 0, 1, 1, 0, 1], [1, 1, 0, 0, 0, 0, 1, 0, 0, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1], ]) # Parameters to define the environment to use params = { 'x_axis_vec': x_axis_sp, 'y_axis_vec': y_axis_sp, 'full_map_obs': False, 'pob': 0, 'max_sensor_dist': 10, 'n_sensors': 36, 'fov': 360, 'normalize_dist_sensors': True, 'n_grid_cells': 0, 'bc_n_ring': 0, 'bc_n_rad': 0, 'bc_dist_rad': 0, 'bc_receptive_field_min': 0, 'bc_receptive_field_max': 0, 'hd_n_cells': 0, 'hd_receptive_field_min': 0, 'hd_receptive_field_max': 0, 'heading': 'circular', 'location': 'none', 'goal_loc': 'none', 'goal_vec': 'none', 'goal_csp': True, 'agent_csp': True, 'goal_csp_egocentric': True, 'csp_dim': dim, } obs_dict = generate_obs_dict(params) csp_offset = map_array.shape[0] / 2 csp_scaling = 5 / (map_array.shape[0] / 2) env = GridWorldEnv( map_array=map_array, # object_locations=object_locations, observations=obs_dict, continuous=True, movement_type='holonomic', dt=0.1, max_steps=1000, fixed_episode_length=False, max_lin_vel=5, max_ang_vel=5, screen_width=300, screen_height=300, csp_scaling=csp_offset, csp_offset=csp_scaling, ) return env
def final(): fig, ax = plt.subplots(5, 4, figsize=(8, 6), tight_layout=True) fig2, ax2 = plt.subplots(5, 3, figsize=(5, 4), tight_layout=True) dim = 256 res = 32 sigma = .15 n_toroid = (dim - 1) // 2 for type_index in range(5): phi_mag_total = np.zeros((n_toroid*3, 1)) for seed in range(3): if type_index == 0: # random SSP rng = np.random.RandomState(seed=seed) X = make_good_unitary(dim, rng=rng) Y = make_good_unitary(dim, rng=rng) md, phi_pos = phi_mag_and_dir(X.v, Y.v) elif type_index == 1: # sub-toroid SSP rng = np.random.RandomState(seed=seed) X, Y = get_fixed_dim_sub_toriod_axes( dim=dim, n_proj=3, scale_ratio=0, scale_start_index=0, rng=rng, eps=0.001, ) md, phi_pos = phi_mag_and_dir(X.v, Y.v) elif type_index == 2: # learned SSP no regularization X, Y = get_axes_from_network( fname='learned_ssp_models/no_reg_model_1layer_1024hs_seed{}.pt'.format(seed), # fname='learned_ssp_models/reg_proper_model_1layer_1024hs_seed{}.pt'.format(seed), # fname='learned_ssp_models/reg_proper_noise_model_1layer_1024hs_seed{}.pt'.format(seed), # fname='learned_ssp_models/reg_proper_phi_decay_model_1layer_1024hs_seed{}.pt'.format(seed), dim=dim ) md, phi_pos = phi_mag_and_dir(X, Y) elif type_index == 3: # learned SSP just weight regularization X, Y = get_axes_from_network( fname='learned_ssp_models/reg_proper_model_1layer_1024hs_seed{}.pt'.format(seed), dim=dim ) md, phi_pos = phi_mag_and_dir(X, Y) elif type_index == 4: # learned SSP weight and phi regularization X, Y = get_axes_from_network( fname='learned_ssp_models/reg_proper_phi_decay_model_1layer_1024hs_seed{}.pt'.format(seed), dim=dim ) md, phi_pos = phi_mag_and_dir(X, Y) gauss_image(phi_pos, ax[type_index, seed], sigma=sigma, res=res) if type_index < 2: X = X.v Y = Y.v im = plot_heatmap(X, Y, np.linspace(-5, 5, 128), np.linspace(-5, 5, 128), ax2[type_index, seed]) phi_mag_total[seed * n_toroid:(seed + 1) * n_toroid, 0] = md[:, 0] if seed == 0: if type_index == 0: # label = 'Fixed SSP' label = 'A ' elif type_index == 1: # label = 'Fixed Grid SSP' label = 'B ' elif type_index == 2: # label = 'Learned SSP' label = 'C ' elif type_index == 3: # label = 'Learned SSP' label = 'D ' elif type_index == 4: # label = 'Learned SSP' label = 'E ' # ax[type_index, seed].set_ylabel( # label, rotation=90, fontsize=18, # # position=(0, .4) # ) ax[type_index, seed].set_ylabel( label, rotation=0, fontsize=18, position=(0, .4) ) ax2[type_index, seed].set_ylabel( label, rotation=0, fontsize=18, position=(0, .4) ) # histogram at the bottom magnitude_histogram(phi_mag_total, ax[type_index, 3]) # if type_index == 0: # # ax[type_index, 1].set_title('Phi Magnitude and Direction') # # ax[type_index, 3].set_title('Phi Magnitude Histogram') # cbar_ax = fig2.add_axes([0.85, 0.05, 0.05, 0.85]) # add_axes([left, bottom, width, height]) cbar_ax = fig2.add_axes([0.85, 0.05, 0.05, 0.90]) fig.colorbar(im, cax=cbar_ax) plt.show()
vec[(i - j) % dim] += mat[i, j] vec /= dim return vec rng = np.random.RandomState(seed=args.seed) xs = np.linspace(args.limit_low, args.limit_high, args.res) axis_vector_type = 'pca' if axis_vector_type == 'standard': dim = 256 X = make_good_unitary(dim, rng=rng) # X = spa.SemanticPointer(dim) # X.make_unitary() X_circ = circulant(X.v) X_vec = circulant_matrix_to_vec(X_circ) # assert (np.all(X_vec == X.v)) elif axis_vector_type == 'covariance': # generate SSP based on a given circulant matrix data = np.load(args.fname) X_circ = data['covariance'] dim = X_circ.shape[0] X_vec = circulant_matrix_to_vec(X_circ) X = spa.SemanticPointer(data=X_vec)
traj_folder_name) if not os.path.exists(traj_path): os.makedirs(traj_path) traj_name = os.path.join(traj_path, 'trajectory_dataset.npz') limit_low = 0 limit_high = args.maze_size if not os.path.exists(base_name): print("Generating base maze data") rng = np.random.RandomState(seed=args.seed) x_axis_sp = make_good_unitary(dim=args.dim, rng=rng) y_axis_sp = make_good_unitary(dim=args.dim, rng=rng) np.random.seed(args.seed) xs_coarse = np.linspace(limit_low, limit_high, args.maze_size) ys_coarse = np.linspace(limit_low, limit_high, args.maze_size) xs = np.linspace(limit_low, limit_high, args.res) ys = np.linspace(limit_low, limit_high, args.res) coarse_mazes = np.zeros((args.n_mazes, args.maze_size, args.maze_size)) fine_mazes = np.zeros((args.n_mazes, args.res, args.res)) solved_mazes = np.zeros( (args.n_mazes, args.n_goals, args.res, args.res, 2)) maze_sps = np.zeros((args.n_mazes, args.dim))
def main(): parser = argparse.ArgumentParser( 'Train a network to clean up a noisy spatial semantic pointer') parser.add_argument('--loss', type=str, default='cosine', choices=['cosine', 'mse']) parser.add_argument('--noise-type', type=str, default='memory', choices=['memory', 'gaussian', 'both']) parser.add_argument( '--sigma', type=float, default=1.0, help='sigma on the gaussian noise if noise-type==gaussian') parser.add_argument('--train-fraction', type=float, default=.8, help='proportion of the dataset to use for training') parser.add_argument( '--n-samples', type=int, default=10000, help= 'Number of memories to generate. Total samples will be n-samples * n-items' ) parser.add_argument('--n-items', type=int, default=12, help='number of items in memory. Proxy for noisiness') parser.add_argument('--dim', type=int, default=512, help='Dimensionality of the semantic pointers') parser.add_argument('--hidden-size', type=int, default=512, help='Hidden size of the cleanup network') parser.add_argument('--limits', type=str, default="-5,5,-5,5", help='The limits of the space') parser.add_argument('--epochs', type=int, default=50) parser.add_argument('--batch-size', type=int, default=32) parser.add_argument('--lr', type=float, default=0.001) parser.add_argument('--momentum', type=float, default=0.9) parser.add_argument('--seed', type=int, default=13) parser.add_argument('--logdir', type=str, default='ssp_cleanup', help='Directory for saved model and tensorboard log') parser.add_argument('--load-model', type=str, default='', help='Optional model to continue training from') parser.add_argument( '--name', type=str, default='', help= 'Name of output folder within logdir. Will use current date and time if blank' ) parser.add_argument('--weight-histogram', action='store_true', help='Save histograms of the weights if set') parser.add_argument('--use-hex-ssp', action='store_true') parser.add_argument('--optimizer', type=str, default='adam', choices=['sgd', 'adam', 'rmsprop']) args = parser.parse_args() args.limits = tuple(float(v) for v in args.limits.split(',')) np.random.seed(args.seed) torch.manual_seed(args.seed) dataset_name = 'data/ssp_cleanup_dataset_dim{}_seed{}_items{}_limit{}_samples{}.npz'.format( args.dim, args.seed, args.n_items, args.limits[1], args.n_samples) final_test_samples = 100 final_test_items = 15 final_test_dataset_name = 'data/ssp_cleanup_test_dataset_dim{}_seed{}_items{}_limit{}_samples{}.npz'.format( args.dim, args.seed, final_test_items, args.limits[1], final_test_samples) if not os.path.exists('data'): os.makedirs('data') rng = np.random.RandomState(seed=args.seed) if args.use_hex_ssp: x_axis_sp, y_axis_sp = get_axes(dim=args.dim, n=3, seed=args.seed) else: x_axis_sp = make_good_unitary(args.dim, rng=rng) y_axis_sp = make_good_unitary(args.dim, rng=rng) if args.noise_type == 'gaussian': # Simple generation clean_ssps = np.zeros((args.n_samples, args.dim)) coords = np.zeros((args.n_samples, 2)) for i in range(args.n_samples): x = np.random.uniform(low=args.limits[0], high=args.limits[1]) y = np.random.uniform(low=args.limits[2], high=args.limits[3]) clean_ssps[i, :] = encode_point(x, y, x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp).v coords[i, 0] = x coords[i, 1] = y # Gaussian noise will be added later noisy_ssps = clean_ssps.copy() else: if os.path.exists(dataset_name): print("Loading dataset") data = np.load(dataset_name) clean_ssps = data['clean_ssps'] noisy_ssps = data['noisy_ssps'] else: print("Generating SSP cleanup dataset") clean_ssps, noisy_ssps, coords = generate_cleanup_dataset( x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp, n_samples=args.n_samples, dim=args.dim, n_items=args.n_items, limits=args.limits, seed=args.seed, ) print("Dataset generation complete. Saving dataset") np.savez( dataset_name, clean_ssps=clean_ssps, noisy_ssps=noisy_ssps, coords=coords, x_axis_vec=x_axis_sp.v, y_axis_vec=x_axis_sp.v, ) # check if the final test set has been generated yet if os.path.exists(final_test_dataset_name): print("Loading final test dataset") final_test_data = np.load(final_test_dataset_name) final_test_clean_ssps = final_test_data['clean_ssps'] final_test_noisy_ssps = final_test_data['noisy_ssps'] else: print("Generating final test dataset") final_test_clean_ssps, final_test_noisy_ssps, final_test_coords = generate_cleanup_dataset( x_axis_sp=x_axis_sp, y_axis_sp=y_axis_sp, n_samples=final_test_samples, dim=args.dim, n_items=final_test_items, limits=args.limits, seed=args.seed, ) print("Final test generation complete. Saving dataset") np.savez( final_test_dataset_name, clean_ssps=final_test_clean_ssps, noisy_ssps=final_test_noisy_ssps, coords=final_test_coords, x_axis_vec=x_axis_sp.v, y_axis_vec=x_axis_sp.v, ) # Add gaussian noise if required if args.noise_type == 'gaussian' or args.noise_type == 'both': noisy_ssps += np.random.normal(loc=0, scale=args.sigma, size=noisy_ssps.shape) n_samples = clean_ssps.shape[0] n_train = int(args.train_fraction * n_samples) n_test = n_samples - n_train assert (n_train > 0 and n_test > 0) train_clean = clean_ssps[:n_train, :] train_noisy = noisy_ssps[:n_train, :] test_clean = clean_ssps[n_train:, :] test_noisy = noisy_ssps[n_train:, :] # NOTE: this dataset is actually generic and can take any input/output mapping dataset_train = CoordDecodeDataset(vectors=train_noisy, coords=train_clean) dataset_test = CoordDecodeDataset(vectors=test_noisy, coords=test_clean) dataset_final_test = CoordDecodeDataset(vectors=final_test_noisy_ssps, coords=final_test_clean_ssps) trainloader = torch.utils.data.DataLoader( dataset_train, batch_size=args.batch_size, shuffle=True, num_workers=0, ) # For testing just do everything in one giant batch testloader = torch.utils.data.DataLoader( dataset_test, batch_size=len(dataset_test), shuffle=False, num_workers=0, ) final_testloader = torch.utils.data.DataLoader( dataset_final_test, batch_size=len(dataset_final_test), shuffle=False, num_workers=0, ) model = FeedForward(dim=dataset_train.dim, hidden_size=args.hidden_size, output_size=dataset_train.dim) # Open a tensorboard writer if a logging directory is given if args.logdir != '': current_time = datetime.now().strftime('%b%d_%H-%M-%S') save_dir = osp.join(args.logdir, current_time) writer = SummaryWriter(log_dir=save_dir) if args.weight_histogram: # Log the initial parameters for name, param in model.named_parameters(): writer.add_histogram('parameters/' + name, param.clone().cpu().data.numpy(), 0) mse_criterion = nn.MSELoss() cosine_criterion = nn.CosineEmbeddingLoss() if args.optimizer == 'sgd': optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) elif args.optimizer == 'rmsprop': optimizer = torch.optim.RMSprop(model.parameters(), lr=args.lr, momentum=args.momentum) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) else: raise NotImplementedError for e in range(args.epochs): print('Epoch: {0}'.format(e + 1)) avg_mse_loss = 0 avg_cosine_loss = 0 n_batches = 0 for i, data in enumerate(trainloader): noisy, clean = data if noisy.size()[0] != args.batch_size: continue # Drop data, not enough for a batch optimizer.zero_grad() outputs = model(noisy) mse_loss = mse_criterion(outputs, clean) # Modified to use CosineEmbeddingLoss cosine_loss = cosine_criterion(outputs, clean, torch.ones(args.batch_size)) avg_cosine_loss += cosine_loss.data.item() avg_mse_loss += mse_loss.data.item() n_batches += 1 if args.loss == 'cosine': cosine_loss.backward() else: mse_loss.backward() # print(loss.data.item()) optimizer.step() print(avg_cosine_loss / n_batches) if args.logdir != '': if n_batches > 0: avg_cosine_loss /= n_batches writer.add_scalar('avg_cosine_loss', avg_cosine_loss, e + 1) writer.add_scalar('avg_mse_loss', avg_mse_loss, e + 1) if args.weight_histogram and (e + 1) % 10 == 0: for name, param in model.named_parameters(): writer.add_histogram('parameters/' + name, param.clone().cpu().data.numpy(), e + 1) print("Testing") with torch.no_grad(): for label, loader in zip(['test', 'final_test'], [testloader, final_testloader]): # Everything is in one batch, so this loop will only happen once for i, data in enumerate(loader): noisy, clean = data outputs = model(noisy) mse_loss = mse_criterion(outputs, clean) # Modified to use CosineEmbeddingLoss cosine_loss = cosine_criterion(outputs, clean, torch.ones(len(loader))) print(cosine_loss.data.item()) if args.logdir != '': # TODO: get a visualization of the performance # show plots of the noisy, clean, and cleaned up with the network # note that the plotting mechanism itself uses nearest neighbors, so has a form of cleanup built in xs = np.linspace(args.limits[0], args.limits[1], 256) ys = np.linspace(args.limits[0], args.limits[1], 256) heatmap_vectors = get_heatmap_vectors(xs, ys, x_axis_sp, y_axis_sp) noisy_coord = ssp_to_loc_v(noisy, heatmap_vectors, xs, ys) pred_coord = ssp_to_loc_v(outputs, heatmap_vectors, xs, ys) clean_coord = ssp_to_loc_v(clean, heatmap_vectors, xs, ys) fig_noisy_coord, ax_noisy_coord = plt.subplots() fig_pred_coord, ax_pred_coord = plt.subplots() fig_clean_coord, ax_clean_coord = plt.subplots() plot_predictions_v(noisy_coord, clean_coord, ax_noisy_coord, min_val=args.limits[0], max_val=args.limits[1], fixed_axes=True) plot_predictions_v(pred_coord, clean_coord, ax_pred_coord, min_val=args.limits[0], max_val=args.limits[1], fixed_axes=True) plot_predictions_v(clean_coord, clean_coord, ax_clean_coord, min_val=args.limits[0], max_val=args.limits[1], fixed_axes=True) writer.add_figure('{}/original_noise'.format(label), fig_noisy_coord) writer.add_figure('{}/test_set_cleanup'.format(label), fig_pred_coord) writer.add_figure('{}/ground_truth'.format(label), fig_clean_coord) # fig_hist = plot_histogram(predictions=outputs, coords=coord) # writer.add_figure('test set histogram', fig_hist) writer.add_scalar('{}/test_cosine_loss'.format(label), cosine_loss.data.item()) writer.add_scalar('{}/test_mse_loss'.format(label), mse_loss.data.item()) # Close tensorboard writer if args.logdir != '': writer.close() torch.save(model.state_dict(), osp.join(save_dir, 'model.pt')) params = vars(args) # # Additionally save the axis vectors used # params['x_axis_vec'] = list(x_axis_sp.v) # params['y_axis_vec'] = list(y_axis_sp.v) with open(osp.join(save_dir, "params.json"), "w") as f: json.dump(params, f)