def forward_pass(device_type): env_name = 'atari_breakout' cfg = default_cfg(algo='APPO', env=env_name) cfg.actor_critic_share_weights = True cfg.hidden_size = 128 cfg.use_rnn = True env = create_env(env_name, cfg=cfg) torch.set_num_threads(1) torch.backends.cudnn.benchmark = True actor_critic = create_actor_critic(cfg, env.observation_space, env.action_space) device = torch.device(device_type) actor_critic.to(device) timing = Timing() with timing.timeit('all'): batch = 128 with timing.add_time('input'): # better avoid hardcoding here... observations = dict( obs=torch.rand([batch, 4, 84, 84]).to(device)) rnn_states = torch.rand([batch, get_hidden_size(cfg)]).to(device) n = 200 for i in range(n): with timing.add_time('forward'): output = actor_critic(observations, rnn_states) log.debug('Progress %d/%d', i, n) log.debug('Timing: %s', timing)
def __init__(self, cfg, num_agents, obs_space, action_space): self.cfg = cfg self.num_agents = num_agents self.envs_per_split = cfg.num_envs_per_worker // cfg.worker_num_splits self.num_traj_buffers = self.calc_num_trajectory_buffers() num_actions = calc_num_actions(action_space) num_action_logits = calc_num_logits(action_space) hidden_size = get_hidden_size(self.cfg) log.debug('Allocating shared memory for trajectories') self.tensors = TensorDict() # policy inputs obs_dict = TensorDict() self.tensors['obs'] = obs_dict if isinstance(obs_space, spaces.Dict): for space_name, space in obs_space.spaces.items(): obs_dict[space_name] = self.init_tensor(space.dtype, space.shape) else: raise Exception('Only Dict observations spaces are supported') # env outputs self.tensors['rewards'] = self.init_tensor(torch.float32, [1]) self.tensors['dones'] = self.init_tensor(torch.bool, [1]) # policy outputs policy_outputs = [ ('actions', num_actions), ('action_logits', num_action_logits), ('log_prob_actions', 1), ('values', 1), ('policy_version', 1), ('rnn_states', hidden_size) ] policy_outputs = [PolicyOutput(*po) for po in policy_outputs] policy_outputs = sorted(policy_outputs, key=lambda policy_output: policy_output.name) for po in policy_outputs: self.tensors[po.name] = self.init_tensor(torch.float32, [po.size]) ensure_memory_shared(self.tensors) # this is for performance optimization # indexing in numpy arrays is faster than in PyTorch tensors self.tensors_individual_transitions = self.tensor_dict_to_numpy(len(self.tensor_dimensions())) self.tensor_trajectories = self.tensor_dict_to_numpy(len(self.tensor_dimensions()) - 1) # create a shared tensor to indicate when the learner is done with the trajectory buffer and # it can be used to store the next trajectory traj_buffer_available_shape = [ self.cfg.num_workers, self.cfg.worker_num_splits, self.envs_per_split, self.num_agents, self.num_traj_buffers, ] self.is_traj_tensor_available = torch.ones(traj_buffer_available_shape, dtype=torch.uint8) self.is_traj_tensor_available.share_memory_() self.is_traj_tensor_available = to_numpy(self.is_traj_tensor_available, 2) # copying small policy outputs (e.g. individual value predictions & action logits) to shared memory is a # bottleneck on the policy worker. For optimization purposes we create additional tensors to hold # just concatenated policy outputs. Rollout workers parse the data and add it to the trajectory buffers # in a proper format policy_outputs_combined_size = sum(po.size for po in policy_outputs) policy_outputs_shape = [ self.cfg.num_workers, self.cfg.worker_num_splits, self.envs_per_split, self.num_agents, policy_outputs_combined_size, ] self.policy_outputs = policy_outputs self.policy_output_tensors = torch.zeros(policy_outputs_shape, dtype=torch.float32) self.policy_output_tensors.share_memory_() self.policy_output_tensors = to_numpy(self.policy_output_tensors, 4) self.policy_versions = torch.zeros([self.cfg.num_policies], dtype=torch.int32) self.policy_versions.share_memory_() # a list of boolean flags to be shared among components that indicate that experience collection should be # temporarily stopped (e.g. due to too much experience accumulated on the learner) self.stop_experience_collection = torch.ones([self.cfg.num_policies], dtype=torch.bool) self.stop_experience_collection.share_memory_()
def enjoy(cfg, max_num_episodes=1000000, max_num_frames=1e9): cfg = load_from_checkpoint(cfg) render_action_repeat = cfg.render_action_repeat if cfg.render_action_repeat is not None else cfg.env_frameskip if render_action_repeat is None: log.warning('Not using action repeat!') render_action_repeat = 1 log.debug('Using action repeat %d during evaluation', render_action_repeat) cfg.env_frameskip = 1 # for evaluation cfg.num_envs = 1 if cfg.record_to: tstamp = datetime.datetime.now().strftime('%Y_%m_%d__%H_%M_%S') cfg.record_to = join(cfg.record_to, f'{cfg.experiment}', tstamp) if not os.path.isdir(cfg.record_to): os.makedirs(cfg.record_to) else: cfg.record_to = None def make_env_func(env_config): return create_env(cfg.env, cfg=cfg, env_config=env_config) env = make_env_func(AttrDict({'worker_index': 0, 'vector_index': 0})) # env.seed(0) is_multiagent = is_multiagent_env(env) if not is_multiagent: env = MultiAgentWrapper(env) if hasattr(env.unwrapped, 'reset_on_init'): # reset call ruins the demo recording for VizDoom env.unwrapped.reset_on_init = False actor_critic = create_actor_critic(cfg, env.observation_space, env.action_space) device = torch.device('cpu' if cfg.device == 'cpu' else 'cuda') actor_critic.model_to_device(device) policy_id = cfg.policy_index checkpoints = LearnerWorker.get_checkpoints( LearnerWorker.checkpoint_dir(cfg, policy_id)) checkpoint_dict = LearnerWorker.load_checkpoint(checkpoints, device) actor_critic.load_state_dict(checkpoint_dict['model']) episode_rewards = [] true_rewards = deque([], maxlen=100) num_frames = 0 last_render_start = time.time() def max_frames_reached(frames): return max_num_frames is not None and frames > max_num_frames obs = env.reset() with torch.no_grad(): for _ in range(max_num_episodes): done = [False] * len(obs) rnn_states = torch.zeros( [env.num_agents, get_hidden_size(cfg)], dtype=torch.float32, device=device) episode_reward = 0 while True: obs_torch = AttrDict(transform_dict_observations(obs)) for key, x in obs_torch.items(): obs_torch[key] = torch.from_numpy(x).to(device).float() policy_outputs = actor_critic(obs_torch, rnn_states, with_action_distribution=True) # sample actions from the distribution by default actions = policy_outputs.actions action_distribution = policy_outputs.action_distribution if isinstance(action_distribution, ContinuousActionDistribution): if not cfg.continuous_actions_sample: # TODO: add similar option for discrete actions actions = action_distribution.means actions = actions.cpu().numpy() rnn_states = policy_outputs.rnn_states for _ in range(render_action_repeat): if not cfg.no_render: target_delay = 1.0 / cfg.fps if cfg.fps > 0 else 0 current_delay = time.time() - last_render_start time_wait = target_delay - current_delay if time_wait > 0: # log.info('Wait time %.3f', time_wait) time.sleep(time_wait) last_render_start = time.time() env.render() obs, rew, done, infos = env.step(actions) episode_reward += np.mean(rew) num_frames += 1 if all(done): true_rewards.append(infos[0].get( 'true_reward', math.nan)) log.info('Episode finished at %d frames', num_frames) if not math.isnan(np.mean(true_rewards)): log.info('true rew %.3f avg true rew %.3f', true_rewards[-1], np.mean(true_rewards)) # VizDoom multiplayer stuff # for player in [1, 2, 3, 4, 5, 6, 7, 8]: # key = f'PLAYER{player}_FRAGCOUNT' # if key in infos[0]: # log.debug('Score for player %d: %r', player, infos[0][key]) break if all(done) or max_frames_reached(num_frames): break if not cfg.no_render: env.render() time.sleep(0.01) episode_rewards.append(episode_reward) last_episodes = episode_rewards[-100:] avg_reward = sum(last_episodes) / len(last_episodes) log.info( 'Episode reward: %f, avg reward for %d episodes: %f', episode_reward, len(last_episodes), avg_reward, ) if max_frames_reached(num_frames): break env.close() return ExperimentStatus.SUCCESS, np.mean(episode_rewards)