From ac642d5d496fe21997d3c14282fe16fe0bf1034b Mon Sep 17 00:00:00 2001 From: clementine Date: Mon, 1 Jul 2024 23:58:20 +0100 Subject: [PATCH] first steps to integrate the trajectory --- neuralplayground/agents/agent_core.py | 44 ++++++++++++++ neuralplayground/debug.py | 15 +++++ neuralplayground/debug_1.py | 85 +++++++++++++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 neuralplayground/debug.py create mode 100644 neuralplayground/debug_1.py diff --git a/neuralplayground/agents/agent_core.py b/neuralplayground/agents/agent_core.py index 983b1d7d..e2aee209 100644 --- a/neuralplayground/agents/agent_core.py +++ b/neuralplayground/agents/agent_core.py @@ -219,6 +219,7 @@ def _act(self, obs): position variation to compute next position """ # Pick direction + direction = super().act(obs) # Normalize direction to step size direction = direction / np.sqrt(np.sum(direction**2)) * self.step_size @@ -260,3 +261,46 @@ def act(self, obs): return self.action_buffer.pop() else: return action + +class TrajectoryGenerator(RandomAgent): + def __init__( + self,step_size: float = 0.02,): + super().__init__(step_size=step_size) + self.action_buffer = [] + self.b = 0.13 * 2 * np.pi # forward velocity rayleigh dist scale (m/sec) + self.sigma = 1 # stdev rotation velocity (rads/sec) + self.mu = 0 # turn angle bias + self.reset() + def reset(self): + self.head_dir = np.random.uniform(0, 2 * np.pi) + self.velocity = 0 + self.turn_angle = 0 + + def act(self, obs, crossed_border): + # Pick direction + """The base model executes a random action from a normal distribution + Parameters + ---------- + obs: + Whatever observation from the environment class needed to choose the right action + Returns + ------- + d_pos: nd.array (2,) + position variation to compute next position + """ + #is_near_wall, turn_angle = self.avoid_wall(position[:, t], head_dir[:, t], room_width, room_depth) + v = np.random.rayleigh(self.b,1) + v[crossed_border] *= 0.25 + + turn_angle = np.zeros(1) + random_turn = np.random.normal(self.mu, self.sigma, 1) + self.turn_angle += self.step_size * random_turn + + self.velocity = self.step_size * v + + d_pos = self.velocity * np.stack([np.cos(self.head_dir), np.sin(self.head_dir)]).reshape(-1) + self.head_dir = self.head_dir + self.turn_angle + + self.head_dir = np.mod(self.head_dir + np.pi, 2 * np.pi) - np.pi + + return d_pos \ No newline at end of file diff --git a/neuralplayground/debug.py b/neuralplayground/debug.py new file mode 100644 index 00000000..0f0bf400 --- /dev/null +++ b/neuralplayground/debug.py @@ -0,0 +1,15 @@ +from neuralplayground.agents import RandomAgent, LevyFlightAgent,TrajectoryGenerator +# Random agent generates a brownian motion. Levy flight is still experimental. + +agent = TrajectoryGenerator() +time_step_size = 0.1 #seg +agent_step_size = 3 + +# Init environment +env = Simple2D(time_step_size = time_step_size, + agent_step_size = agent_step_size, + arena_x_limits=(-100, 100), + arena_y_limits=(-100, 100)) + + +print('testing stuff ') \ No newline at end of file diff --git a/neuralplayground/debug_1.py b/neuralplayground/debug_1.py new file mode 100644 index 00000000..42067429 --- /dev/null +++ b/neuralplayground/debug_1.py @@ -0,0 +1,85 @@ +import os +import matplotlib.pyplot as plt +os.environ['KMP_DUPLICATE_LIB_OK']='True' +from neuralplayground.arenas import Simple2D +from neuralplayground.agents import RandomAgent, LevyFlightAgent, TrajectoryGenerator +from neuralplayground.experiments import Sargolini2006Data +from neuralplayground.agents import TrajectoryGenerator, Burak2009, Sorscher2022exercise, SorscherIdealRNN +from neuralplayground.utils import PlaceCells, get_2d_sort +from neuralplayground.plotting import plot_trajectory_place_cells_activity, plot_ratemaps, compute_ratemaps +from neuralplayground.config import load_plot_config +load_plot_config() +import matplotlib.pyplot as plt +import numpy as np +import torch +from tqdm import tqdm +from mpl_toolkits.mplot3d import Axes3D +import os +import pickle +import matplotlib as mpl +import scipy +np.random.seed(0) + +# Use GPU if available +device = torch.device("cpu") +agent_step_size = 0.02 +agent = TrajectoryGenerator(step_size = agent_step_size) +time_step_size = 0.01 + + +# Init environment +env = Simple2D(time_step_size = time_step_size, + agent_step_size = agent_step_size, + arena_x_limits=(-2, 2), +arena_y_limits=(-2, 2)) + +n_steps = 5000#50000 + +# Initialize environment +obs, state, crossed = env.reset() +for i in range(n_steps): + # Observe to choose an action + action = agent.act(obs,crossed) + # Run environment for given action + obs, state, reward, crossed = env.step(action) +ax = env.plot_trajectory() +ax.grid() +plt.show() + + +print('testing stuff ') + + +# We'll use a longer sequence just for plotting purposes +# Training will be done with short sequences +sequence_length = 300 +batch_size = 4 +room_width = 2.2 +room_depth = 2.2 +# Arena dimensions Just 2D +room_width = 2.2 +room_depth = 2.2 + +# We'll use a longer sequence just for plotting purposes +# Training will be done with short sequences +sequence_length = 300 +batch_size = 4 + +# Place cells parameters +n_place_cells = 512 +place_cell_rf = 0.12 +surround_scale = 2.0 +periodic = False +difference_of_gaussians = True +place_cells = PlaceCells(Np=n_place_cells, + place_cell_rf=place_cell_rf, + surround_scale=surround_scale, + room_width=room_width, + room_depth=room_depth, + periodic=periodic, + DoG=difference_of_gaussians, + device=device) +device = torch.device("cpu") +generator = TrajectoryGenerator(sequence_length, batch_size, room_width, room_depth, device, place_cells=place_cells) +traj = generator.generate_trajectory(room_width, room_depth, batch_size) +x, y = traj["target_x"], traj["target_y"] \ No newline at end of file