diff --git a/agent/A2C_agent.py b/agent/A2C_agent.py index 41fca9f..b80c1ad 100644 --- a/agent/A2C_agent.py +++ b/agent/A2C_agent.py @@ -25,11 +25,8 @@ class A2CAgent: self.total_steps = 0 self.states = self.task.reset() - self.episode_counts = np.zeros(config.num_workers) self.episode_rewards = np.zeros(config.num_workers) - self.total_rewards = np.zeros(config.num_workers) - self.prev_episode_counts = 0.0 - self.prev_total_rewards = 0.0 + self.last_episode_rewards = np.zeros(config.num_workers) def close(self): self.task.close() @@ -56,13 +53,9 @@ class A2CAgent: config = self.config for _ in range(config.iteration_log_interval): self.iteration(deterministic) - new_episode_counts = np.sum(self.episode_counts) - new_total_rewards = np.sum(self.total_rewards) - avg_reward = (new_total_rewards - self.prev_total_rewards) / \ - (new_episode_counts - self.prev_episode_counts + 1e-5) - self.prev_total_rewards = new_total_rewards - self.prev_episode_counts = new_episode_counts - return avg_reward, config.rollout_length * config.num_workers * \ + config.logger.info('max/min reward %f/%f' % + (np.max(self.last_episode_rewards), np.min(self.last_episode_rewards))) + return self.last_episode_rewards.mean(), config.rollout_length * config.num_workers * \ config.iteration_log_interval def iteration(self, deterministic=False): @@ -82,8 +75,7 @@ class A2CAgent: for i, terminal in enumerate(terminals): if terminals[i]: next_states[i] = self.task.reset(i) - self.episode_counts[i] += 1 - self.total_rewards[i] += self.episode_rewards[i] + self.last_episode_rewards[i] = self.episode_rewards[i] self.episode_rewards[i] = 0 rollout.append([prob, log_prob, value, actions, rewards, 1 - terminals]) @@ -112,11 +104,16 @@ class A2CAgent: prob, log_prob, value, actions, returns, advantages = map(lambda x: torch.cat(x, dim=0), zip(*processed_rollout)) policy_loss = -log_prob.gather(1, Variable(actions)) * Variable(advantages) - policy_loss += config.entropy_weight * torch.sum(prob * log_prob, dim=1, keepdim=True) - value_loss = config.value_loss_weight * 0.5 * (Variable(returns) - value).pow(2) + entropy_loss = torch.sum(prob * log_prob, dim=1, keepdim=True) + value_loss = 0.5 * (Variable(returns) - value).pow(2) + + self.config.logger.scalar_summary('policy_loss', np.mean(policy_loss.data.cpu().numpy())) + self.config.logger.scalar_summary('entropy_loss', np.mean(entropy_loss.data.cpu().numpy())) + self.config.logger.scalar_summary('value_loss', np.mean(value_loss.data.cpu().numpy())) self.optimizer.zero_grad() - (policy_loss + value_loss).mean().backward() + (policy_loss + config.entropy_weight * entropy_loss + + config.value_loss_weight * value_loss).mean().backward() nn.utils.clip_grad_norm(self.network.parameters(), config.gradient_clip) self.optimizer.step() diff --git a/component/atari_wrapper.py b/component/atari_wrapper.py index b100eda..45f938d 100644 --- a/component/atari_wrapper.py +++ b/component/atari_wrapper.py @@ -1,11 +1,8 @@ -# This file is apdated from -# https://raw.githubusercontent.com/transedward/pytorch-dqn/master/utils/atari_wrapper.py - import numpy as np from collections import deque import gym from gym import spaces -from PIL import Image +from skimage import color, transform class NoopResetEnv(gym.Wrapper): def __init__(self, env=None, noop_max=30): @@ -105,15 +102,6 @@ class MaxAndSkipEnv(gym.Wrapper): self._obs_buffer.append(obs) return obs -def _process_frame84_rgb(frame): - img = np.reshape(frame, [210, 160, 3]).astype(np.float32) - img = Image.fromarray(img) - resized_screen = img.resize((84, 110, 3), Image.BILINEAR) - resized_screen = np.array(resized_screen) - x_t = resized_screen[18:102, :, :] - x_t = x_t.reshape((84, 84, 3)) - return x_t - class DatasetEnv(gym.Wrapper): def __init__(self, env=None): super(DatasetEnv, self).__init__(env) @@ -138,43 +126,24 @@ class DatasetEnv(gym.Wrapper): self.saved_obs.append(obs) return obs -def _process_frame84(frame): - img = np.reshape(frame, [210, 160, 3]).astype(np.float32) - img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114 - img = Image.fromarray(img) - resized_screen = img.resize((84, 110), Image.BILINEAR) - resized_screen = np.array(resized_screen) - x_t = resized_screen[18:102, :] - x_t = np.reshape(x_t, [1, 84, 84]) - return x_t.astype(np.uint8) - -def _process_frame42(frame): - img = np.reshape(frame, [210, 160, 3]).astype(np.float32) - img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114 - img = img[34:34 + 160, :160] - img = Image.fromarray(img) - img = img.resize((80, 80), Image.BILINEAR) - img = img.resize((42, 42), Image.BILINEAR) - resized_screen = np.array(img).reshape(1, 42, 42) - return resized_screen.astype(np.uint8) - class ProcessFrame(gym.Wrapper): def __init__(self, env=None, frame_size=84): super(ProcessFrame, self).__init__(env) + self.frame_size = frame_size self.observation_space = spaces.Box(low=0, high=255, shape=(1, frame_size, frame_size)) - if frame_size == 84: - self.process_fn = _process_frame84 - elif frame_size == 42: - self.process_fn = _process_frame42 - else: - assert False, "Unknown frame size" + + def process(self, obs): + obs = color.rgb2gray(obs) + obs = transform.resize(obs, (self.frame_size, self.frame_size), mode='constant') + obs = (255 * obs).astype(np.uint8).reshape((1, ) + obs.shape) + return obs def _step(self, action): obs, reward, done, info = self.env.step(action) - return self.process_fn(obs), reward, done, info + return self.process(obs), reward, done, info def _reset(self): - return self.process_fn(self.env.reset()) + return self.process(self.env.reset()) class NormalizeFrame(gym.Wrapper): def __init__(self, env=None): diff --git a/main.py b/main.py index 7074730..947d3ed 100644 --- a/main.py +++ b/main.py @@ -168,22 +168,22 @@ def a2c_pixel_atari(name): history_length=config.history_length) config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers) task = config.task_fn() - config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.0007) - # config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.0001) + # config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.0007) + config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.0001) # config.network_fn = lambda: OpenAIActorCriticConvNet( config.network_fn = lambda: NatureActorCriticConvNet( - config.history_length, task.task.env.action_space.n, gpu=0) + config.history_length, task.task.env.action_space.n, gpu=2) config.reward_shift_fn = lambda r: np.sign(r) config.policy_fn = SamplePolicy config.discount = 0.99 - config.no_gae = True + config.no_gae = False config.gae_tau = 0.97 config.entropy_weight = 0.01 config.rollout_length = 5 config.test_interval = 0 config.iteration_log_interval = 100 config.gradient_clip = 0.5 - config.logger = Logger('./log', logger) + config.logger = Logger('./log', logger, skip=True) run_episodes(A2CAgent(config)) def a3c_continuous(): diff --git a/network/conv_network.py b/network/conv_network.py index 15fbb08..2f672f3 100644 --- a/network/conv_network.py +++ b/network/conv_network.py @@ -126,8 +126,8 @@ class NatureActorCriticConvNet(nn.Module, ActorCriticNet): super(NatureActorCriticConvNet, self).__init__() self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) - self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 64, 512) + self.conv3 = nn.Conv2d(64, 32, kernel_size=3, stride=1) + self.fc4 = nn.Linear(7 * 7 * 32, 512) self.fc_actor = nn.Linear(512, n_actions) self.fc_critic = nn.Linear(512, 1) diff --git a/utils/tf_logger.py b/utils/tf_logger.py index 72e2905..6c2c79b 100644 --- a/utils/tf_logger.py +++ b/utils/tf_logger.py @@ -5,21 +5,35 @@ ####################################################################### from tensorboardX import SummaryWriter +import os +import numpy as np class Logger(object): def __init__(self, log_dir, vanilla_logger, skip=False): - self.writer = SummaryWriter(log_dir) + for f in os.listdir(log_dir): + os.remove('%s/%s' % (log_dir, f)) + if not skip: + self.writer = SummaryWriter(log_dir) self.info = vanilla_logger.info self.debug = vanilla_logger.debug self.warning = vanilla_logger.warning self.skip = skip + self.step = 0 - def scalar_summary(self, tag, value, step): + def scalar_summary(self, tag, value, step=None): if self.skip: return + if step is None: + step = self.step + self.step += 1 + if np.isscalar(value): + value = np.asarray([value]) self.writer.add_scalar(tag, value, step) - def histo_summary(self, tag, values, step): + def histo_summary(self, tag, values, step=None): if self.skip: return + if step is None: + step = self.step + self.step += 1 self.writer.add_histogram(tag, values, step, bins=1000) \ No newline at end of file