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)