def main(): N = 12 COMM_RANGE = 2.5 headless = False leader = True dir_name = "data" try: os.mkdir(dir_name) except FileExistsError: pass data_name = 'flocking_N=%d' % N data_path = os.path.join(dir_name, '%s.pt' % data_name) data = Trainer(data_path=data_path) z = Uniform(low=2.0 * torch.ones(N, 1), high=5.0 * torch.ones(N, 1)) # uniform between 0 and 1 in the z direction xy_normal = Normal(torch.zeros(N, 2), 1.25) # gaussian in the xy direction dist = CombinedDistribution([xy_normal, z], mixer='cat', dim=1) model = Reynolds(N=N, D=6, K=1, OUT_DIM=3) env = gym.make('mrs-v0', state_fn=state_fn, update_fn=update_fn, done_fn=done_fn, N_AGENTS=N, START_POS=dist, K_HOPS=1, COMM_RANGE=COMM_RANGE, ACTION_TYPE='set_target_vel', HEADLESS=headless) if leader: leader_action_policy = RandomAction() action_fn = leader_action_policy.action_fn environment = env.get_env() leader_agent = environment.agents[0] environment.set_colour(leader_agent, [1., 0., 0.]) else: action_fn = lambda action, state: action data.save_trainer_onexit() data = generate_mrs(env=env, model=model, action_fn=action_fn, trainer=data, datapoints=1000, episode_length=200)
def main(): # Parameters N = 12 D = 6 K = 1 COMM_RANGE = 2.5 datapoints = 1000 episode_length = 200 headless = True leader = True # File Paths dir_path = "data" evaldata_path = os.path.join(dir_path, "%s_data.pt") try: os.mkdir(dir_path) except FileExistsError: pass # Initialise Models reynolds = Reynolds(N, D, 1, 3) models = { "reynolds": reynolds, "random": RandomController(OUT_DIM=3) } # we will compare reynolds flocking to a random model # Create Environment z = Uniform(low=2.0 * torch.ones(N, 1), high=5.0 * torch.ones(N, 1)) xy_normal = Normal(torch.zeros(N, 2), 1.0) dist = CombinedDistribution( [xy_normal, z], mixer='cat', dim=1) # create custom starting state distribution env = gym.make('mrs-v0', state_fn=state_fn, update_fn=update_fn, N_AGENTS=N, START_POS=dist, K_HOPS=1, COMM_RANGE=COMM_RANGE, ACTION_TYPE='set_target_vel', HEADLESS=headless) startpos = [ env.generate_start_pos() for _ in range(int(datapoints / episode_length * 2)) ] env.START_POS = StartPosGenerator( startpos) # use the same starting state for each model # Generate and Analyse Data analysers = {} for name, model in models.items(): print(name) data = Trainer(K=K) is_data_loaded = data.load_trainer( path=evaldata_path % name) # load simulation data if it exists if not is_data_loaded: # generate data if it does not exist data.save_trainer_onexit(path=evaldata_path % name) simulate(env=env, model=model, trainer=data, datapoints=datapoints, episode_length=episode_length, leader=leader) analysers[name] = MRSAnalytics( data ) # compute flocking metrics (separation, cohesion, leader dist) analysers[name].name = name # Draw Plots plot_separation(*analysers.values()) plot_cohesion(*analysers.values()) plot_leader_dist(*analysers.values()) # Show show_plots()