def __init__(self, env_config, models: Dict[Tuple, BaseNeuralNet], lr=0.0003, gamma=0.998, clip_param=0.2, entropy_coef=0.025, value_loss_coef=0.05, device='cuda'): self.memory = dict() self.device = device self.env_config = env_config env = foundation.make_env_instance(**env_config) self.models = dict() obs = env.reset() for key, value in models.items(): self.models[key] = value(device=self.device) self.models[key].build_models(ObservationBatch(obs, key)) self.memory[key] = Memory() self.lr = lr self.gamma = gamma self.clip_param = clip_param self.entropy_coef = entropy_coef self.value_loss_coef = value_loss_coef self.optimizers = dict() for key, value in self.models.items(): self.optimizers[key] = torch.optim.Adam(value.parameters(), lr=self.lr)
def __init__(self, env_config, force_dense_logging: bool = False): self.env: BaseEnvironment = foundation.make_env_instance(**{ **ENV_CONF_DEFAULT, **env_config, }) self.observation_space = OBS_SPACE_AGENT self.action_space = ACT_SPACE_AGENT self.force_dense_logging = force_dense_logging
def __init__(self, env_config, verbose=False): self.env_config_dict = env_config["env_config_dict"] # Adding env id in the case of multiple environments if hasattr(env_config, "worker_index"): self.env_id = ( env_config["num_envs_per_worker"] * (env_config.worker_index - 1) ) + env_config.vector_index else: self.env_id = None self.env = foundation.make_env_instance(**self.env_config_dict) self.verbose = verbose self.sample_agent_idx = str(self.env.all_agents[0].idx) obs = self.env.reset() self.observation_space = self._dict_to_spaces_dict(obs["0"]) self.observation_space_pl = self._dict_to_spaces_dict(obs["p"]) if self.env.world.agents[0].multi_action_mode: self.action_space = spaces.MultiDiscrete( self.env.get_agent(self.sample_agent_idx).action_spaces ) self.action_space.dtype = np.int64 self.action_space.nvec = self.action_space.nvec.astype(np.int64) else: self.action_space = spaces.Discrete( self.env.get_agent(self.sample_agent_idx).action_spaces ) self.action_space.dtype = np.int64 if self.env.world.planner.multi_action_mode: self.action_space_pl = spaces.MultiDiscrete( self.env.get_agent("p").action_spaces ) self.action_space_pl.dtype = np.int64 self.action_space_pl.nvec = self.action_space_pl.nvec.astype(np.int64) else: self.action_space_pl = spaces.Discrete( self.env.get_agent("p").action_spaces ) self.action_space_pl.dtype = np.int64 self._seed = None if self.verbose: print("[EnvWrapper] Spaces") print("[EnvWrapper] Obs (a) ") pretty_print(self.observation_space) print("[EnvWrapper] Obs (p) ") pretty_print(self.observation_space_pl) print("[EnvWrapper] Action (a)", self.action_space) print("[EnvWrapper] Action (p)", self.action_space_pl)
def rollout(self, key, n_rollouts, num_steps): env = foundation.make_env_instance(**self.env_config) t = trange(n_rollouts, desc='Rollout') for rollout in t: obs = env.reset() states, actions, logprobs, rewards, values, done, hcs = [], [], [], [], [], [], [] hc = (torch.zeros(1, len(key), self.models[key].lstm_size, device=self.device), torch.zeros(1, len(key), self.models[key].lstm_size, device=self.device)) for step in range(num_steps): obs_batch = ObservationBatch(obs, key) hcs.append(hc) dist, value, hc = self.models[key](obs_batch, hc) hc = (hc[0].detach(), hc[1].detach()) a = dist.sample().detach() value = value.squeeze() logprob = dist.log_prob(a).detach() action_dict = dict( (i, a.detach().cpu().numpy()) for i, a in zip(obs_batch.order, a.argmax(-1))) next_obs, rew, is_done, info = env.step(action_dict) states.append(obs_batch) actions.append(a.argmax(-1).detach()) logprobs.append(logprob) rewards.append(np.array([rew[k] for k in key])) values.append(value) done.append(is_done['__all__']) obs = next_obs obs_batch = ObservationBatch(obs, key) _, next_value, hc = self.models[key](obs_batch, hc) next_value = next_value.detach().cpu().numpy() values = torch.stack(values).detach().cpu().numpy() discounted_rewards = self.compute_gae(next_value, rewards, done, values) advantage = (discounted_rewards - values).tolist() discounted_rewards = discounted_rewards.tolist() self.memory[key].add_trace(states, actions, logprobs, discounted_rewards, advantage, hcs)
def set_env_config(self): """Set up a sample environment config""" self.env_config = { # ===== STANDARD ARGUMENTS ====== "n_agents": 4, # Number of non-planner agents "world_size": [15, 15], # [Height, Width] of the env world "episode_length": 1000, # Number of time-steps per episode # In multi-action-mode, the policy selects an action for each action # subspace (defined in component code) # Otherwise, the policy selects only 1 action "multi_action_mode_agents": False, "multi_action_mode_planner": True, # When flattening observations, concatenate scalar & vector observations # before output # Otherwise, return observations with minimal processing "flatten_observations": False, # When Flattening masks, concatenate each action subspace mask # into a single array # Note: flatten_masks = True is recommended for masking action logits "flatten_masks": True, # ===== COMPONENTS ===== # Which components to use "components": [ # (1) Building houses {"Build": {}}, # (2) Trading collectible resources {"ContinuousDoubleAuction": {"max_num_orders": 5}}, # (3) Movement and resource collection {"Gather": {}}, ], # ===== SCENARIO ===== # Which scenario class to use "scenario_name": "uniform/simple_wood_and_stone", # (optional) kwargs of the chosen scenario class "starting_agent_coin": 10, "starting_stone_coverage": 0.10, "starting_wood_coverage": 0.10, } # Create an environment instance from the config self.env = foundation.make_env_instance(**self.env_config)
{'Gather': {}}, ], # ===== SCENARIO ===== # Which scenario class to use (specified by the class's name in the Scenario Registry) 'scenario_name': 'uniform/simple_wood_and_stone', # (optional) kwargs of the chosen scenario class 'starting_agent_coin': 10, 'starting_stone_coverage': 0.10, 'starting_wood_coverage': 0.10, } env = foundation.make_env_instance(**env_config) obs = env.reset() uniform_cls = scenario_registry.get(env_config['scenario_name']) isinstance(env, uniform_cls) isinstance(env, BaseEnvironment) @scenario_registry.add class EmptyScenario(BaseEnvironment): name = "Empty" required_entities = [] def reset_layout(self): """Resets the state of the world object (self.world).""" pass
def main(): log_dir = Path(__file__).parent / f"exp{int(time.time())}" log_dir.mkdir() env = foundation.make_env_instance(**env_config) state = env.reset() if random_seed: torch.manual_seed(random_seed) env.seed(random_seed) memory = [Memory() for _ in range(env.n_agents)] action_dim = state['0'][ "action_mask"].size # todo mask tells which action cannot be taken state_dim = state['0']["flat"].size ppo = PPO(state_dim, action_dim, n_latent_var, lr, betas, gamma, K_epochs, eps_clip, device) # logging variables running_reward = 0 time_step = 0 # training loop for i_episode in range(1, max_episodes + 1): if i_episode % graphical_episode_log_frequency == 0: obs = env.reset(force_dense_logging=True) else: obs = env.reset() for t in range(env.episode_length): time_step += 1 actions = sample_random_actions( env, obs ) # Initialize dict with random actions then fill with selected values for agent_id in range(env.n_agents): agent_id_str = str(agent_id) memory_agent = memory[agent_id] action_mask = torch.tensor(state[agent_id_str]["action_mask"], device=device) agent_state = state[agent_id_str]["flat"] action = ppo.policy_old.act(agent_state, memory_agent, action_mask) actions[agent_id_str] = action state, reward, done, info = env.step(actions) # Saving reward and is_terminals: for agent_id in range(env.n_agents): agent_id_str = str(agent_id) memory_agent = memory[agent_id] agent_reward = -reward[agent_id_str] memory_agent.rewards.append(agent_reward) memory_agent.is_terminals.append(done) running_reward += agent_reward # update if its time if time_step % update_timestep == 0: ppo.update(memory_agent) memory_agent.clear_memory() time_step = 0 if done['__all__']: break # save every 500 episodes if i_episode % 100 == 0: torch.save(ppo.policy.state_dict(), log_dir / f'ckpt-{i_episode}.pth') # logging if i_episode % log_interval == 0: running_reward = float((running_reward / log_interval)) print(f'Episode {i_episode} \t Avg reward: {running_reward}') running_reward = 0 if i_episode % graphical_episode_log_frequency == 0: dense_log = env.previous_episode_dense_log (fig0, fig1, fig2), incomes, endows, c_trades, all_builds = plotting.breakdown( dense_log) print(f"{incomes=}, {endows=}, {c_trades=}, {all_builds=}") # fig0.savefig(log_dir / f"fig0-{i_episode}.png", dpi=fig0.dpi) fig1.savefig(log_dir / f"fig1-{i_episode:04d}.png", dpi=fig1.dpi) fig2.savefig(log_dir / f"fig2-{i_episode:04d}.png", dpi=fig2.dpi) plt.close(fig0) plt.close(fig1) plt.close(fig2) with open(log_dir / f'logs-{i_episode:04d}.pickle', 'wb') as handle: pickle.dump(dense_log, handle, protocol=pickle.HIGHEST_PROTOCOL)