diff --git a/main.py b/main.py index fff949e..3fa2517 100644 --- a/main.py +++ b/main.py @@ -355,6 +355,29 @@ def categorical_dqn_cart_pole(): config.categorical_n_atoms = 50 run_episodes(CategoricalDQNAgent(config)) +def categorical_dqn_pixel_atari(name): + config = Config() + config.history_length = 4 + config.task_fn = lambda: PixelAtari(name, no_op=30, frame_skip=4, normalized_state=False, + history_length=config.history_length) + action_dim = config.task_fn().action_dim + config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.00025, eps=0.01 / 32) + config.network_fn = lambda: CategoricalConvNet(config.history_length, action_dim, config.categorical_n_atoms, gpu=0) + config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) + config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) + config.reward_shift_fn = lambda r: np.sign(r) + config.discount = 0.99 + config.target_network_update_freq = 10000 + config.exploration_steps= 50000 + config.logger = Logger('./log', logger) + config.test_interval = 10 + config.test_repetitions = 1 + config.double_q = False + config.categorical_v_max = 10 + config.categorical_v_min = -10 + config.categorical_n_atoms = 51 + run_episodes(CategoricalDQNAgent(config)) + if __name__ == '__main__': mkdir('data') mkdir('data/video') @@ -364,7 +387,7 @@ if __name__ == '__main__': logger.setLevel(logging.INFO) # dqn_cart_pole() - categorical_dqn_cart_pole() + # categorical_dqn_cart_pole() # async_cart_pole() # a3c_cart_pole() # a2c_cart_pole() @@ -374,6 +397,7 @@ if __name__ == '__main__': # ddpg_continuous() # dqn_pixel_atari('PongNoFrameskip-v4') + categorical_dqn_pixel_atari('PongNoFrameskip-v4') # async_pixel_atari('PongNoFrameskip-v4') # a3c_pixel_atari('PongNoFrameskip-v4') # a2c_pixel_atari('PongNoFrameskip-v4') diff --git a/network/conv_network.py b/network/conv_network.py index 2f672f3..31be959 100644 --- a/network/conv_network.py +++ b/network/conv_network.py @@ -141,3 +141,24 @@ class NatureActorCriticConvNet(nn.Module, ActorCriticNet): x = x.view(x.size(0), -1) phi = F.relu(self.fc4(x)) return phi + +class CategoricalConvNet(nn.Module, CategoricalNet): + def __init__(self, in_channels, n_actions, n_atoms, gpu=0): + super(CategoricalConvNet, 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.fc_categorical = nn.Linear(512, n_actions * n_atoms) + self.n_actions = n_actions + self.n_atoms = n_atoms + BasicNet.__init__(self, gpu) + + def forward(self, x): + x = self.variable(x) + y = F.relu(self.conv1(x)) + y = F.relu(self.conv2(y)) + y = F.relu(self.conv3(y)) + y = y.view(y.size(0), -1) + y = F.relu(self.fc4(y)) + return y \ No newline at end of file