From 015c7753ea3fd7a6f461c3e22e7fd0d8312a1958 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Tue, 10 Apr 2018 23:15:33 -0600 Subject: [PATCH] Continuous actor critic wrapper --- agent/PPO_agent.py | 47 +++++++++++++---------------------------- main.py | 18 ++++++++++------ network/base_network.py | 41 +++++++++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 39 deletions(-) diff --git a/agent/PPO_agent.py b/agent/PPO_agent.py index 98b724f..bc87326 100644 --- a/agent/PPO_agent.py +++ b/agent/PPO_agent.py @@ -19,12 +19,7 @@ class PPOAgent(BaseAgent): BaseAgent.__init__(self) self.config = config self.task = config.task_fn() - self.network = DisjointActorCriticNet(self.task.state_dim, self.task.action_dim, - config.actor_network_fn, config.critic_network_fn) - self.actor = self.network.actor - self.critic = self.network.critic - self.actor_opt = config.actor_optimizer_fn(self.actor.parameters()) - self.critic_opt = config.critic_optimizer_fn(self.critic.parameters()) + self.network = config.network_fn(self.task.state_dim, self.task.action_dim) self.total_steps = 0 self.episode_rewards = np.zeros(config.num_workers) self.last_episode_rewards = np.zeros(config.num_workers) @@ -36,12 +31,7 @@ class PPOAgent(BaseAgent): rollout = [] states = self.states for i in range(config.rollout_length): - mean, std, log_std = self.actor.predict(states) - values = self.critic.predict(states) - dist = torch.distributions.Normal(mean, std) - actions = dist.sample() - log_probs = dist.log_prob(actions).detach() - log_probs = torch.sum(log_probs, dim=1, keepdim=True) + actions, log_probs, _, values = self.network.predict(states) next_states, rewards, terminals, _ = self.task.step(actions.data.cpu().numpy()) self.episode_rewards += rewards rewards = config.reward_normalizer(rewards) @@ -50,22 +40,22 @@ class PPOAgent(BaseAgent): self.last_episode_rewards[i] = self.episode_rewards[i] self.episode_rewards[i] = 0 next_states = config.state_normalizer(next_states) - rollout.append([states, values, actions, log_probs, rewards, 1 - terminals]) + rollout.append([states, values.detach(), actions.detach(), log_probs.detach(), rewards, 1 - terminals]) states = next_states self.states = states - pending_value = self.critic.predict(states) + pending_value = self.network.predict(states)[-1] rollout.append([states, pending_value, None, None, None, None]) processed_rollout = [None] * (len(rollout) - 1) - advantages = self.actor.tensor(np.zeros((config.num_workers, 1))) + advantages = self.network.tensor(np.zeros((config.num_workers, 1))) returns = pending_value.data for i in reversed(range(len(rollout) - 1)): states, value, actions, log_probs, rewards, terminals = rollout[i] - terminals = self.actor.tensor(terminals).unsqueeze(1) - rewards = self.actor.tensor(rewards).unsqueeze(1) - actions = self.actor.variable(actions) - states = self.actor.variable(states) + terminals = self.network.tensor(terminals).unsqueeze(1) + rewards = self.network.tensor(rewards).unsqueeze(1) + actions = self.network.variable(actions) + states = self.network.variable(states) next_value = rollout[i + 1][1] returns = rewards + config.discount * terminals * returns if not config.use_gae: @@ -85,34 +75,27 @@ class PPOAgent(BaseAgent): batcher.shuffle() while not batcher.end(): batch_indices = batcher.next_batch()[0] - batch_indices = self.actor.variable(batch_indices, torch.LongTensor) + batch_indices = self.network.variable(batch_indices, torch.LongTensor) sampled_states = states[batch_indices] sampled_actions = actions[batch_indices] sampled_log_probs_old = log_probs_old[batch_indices] sampled_returns = returns[batch_indices] sampled_advantages = advantages[batch_indices] - mean, std, log_std = self.actor.predict(sampled_states) - dist = torch.distributions.Normal(mean, std) - log_probs = dist.log_prob(sampled_actions) - log_probs = torch.sum(log_probs, dim=1, keepdim=True) + _, log_probs, _, values = self.network.predict(sampled_states, sampled_actions) ratio = (log_probs - sampled_log_probs_old).exp() obj = ratio * sampled_advantages obj_clipped = ratio.clamp(1.0 - self.config.ppo_ratio_clip, 1.0 + self.config.ppo_ratio_clip) * sampled_advantages policy_loss = -torch.min(obj, obj_clipped).mean(0) - v = self.critic.predict(sampled_states) - value_loss = 0.5 * (sampled_returns - v).pow(2).mean() + value_loss = 0.5 * (sampled_returns - values).pow(2).mean() - self.actor_opt.zero_grad() - self.critic_opt.zero_grad() + self.network.zero_grad() policy_loss.backward() value_loss.backward() - nn.utils.clip_grad_norm(self.actor.parameters(), config.gradient_clip) - nn.utils.clip_grad_norm(self.critic.parameters(), config.gradient_clip) - self.actor_opt.step() - self.critic_opt.step() + nn.utils.clip_grad_norm(self.network.parameters(), config.gradient_clip) + self.network.step() steps = config.rollout_length * config.num_workers self.total_steps += steps diff --git a/main.py b/main.py index 9ca98d0..354baa3 100644 --- a/main.py +++ b/main.py @@ -224,16 +224,20 @@ def ppo_continuous(): config.num_workers = 1 # task_fn = lambda log_dir: Pendulum(log_dir=log_dir) # task_fn = lambda log_dir: Roboschool('RoboschoolInvertedPendulum-v1', log_dir=log_dir) - # task_fn = lambda log_dir: Roboschool('RoboschoolAnt-v1', log_dir=log_dir) + task_fn = lambda log_dir: Roboschool('RoboschoolAnt-v1', log_dir=log_dir) # task_fn = lambda log_dir: Roboschool('RoboschoolReacher-v1', log_dir=log_dir) - task_fn = lambda log_dir: Roboschool('RoboschoolHopper-v1', log_dir=log_dir) + # task_fn = lambda log_dir: Roboschool('RoboschoolHopper-v1', log_dir=log_dir) # task_fn = lambda log_dir: DMControl('cartpole', 'balance', log_dir=log_dir) # task_fn = lambda log_dir: DMControl('hopper', 'hop', log_dir=log_dir) config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, log_dir=get_default_log_dir(ppo_continuous.__name__)) - config.actor_network_fn = lambda state_dim, action_dim: GaussianActorNet(state_dim, action_dim) - config.critic_network_fn = lambda state_dim, action_dim: GaussianCriticNet(state_dim) - config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, 3e-4, eps=1e-5) - config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, 3e-4, eps=1e-5) + actor_network_fn = lambda state_dim, action_dim: GaussianActorNet(state_dim, action_dim) + critic_network_fn = lambda state_dim: GaussianCriticNet(state_dim) + actor_optimizer_fn = lambda params: torch.optim.Adam(params, 3e-4, eps=1e-5) + critic_optimizer_fn = lambda params: torch.optim.Adam(params, 3e-4, eps=1e-5) + config.network_fn = lambda state_dim, action_dim: \ + ContinuousActorCriticWrapper(state_dim, action_dim, actor_network_fn, + critic_network_fn, actor_optimizer_fn, + critic_optimizer_fn) # config.state_normalizer = RunningStatsNormalizer() config.discount = 0.99 config.use_gae = True @@ -332,7 +336,7 @@ if __name__ == '__main__': # dqn_ram_atari('Breakout-ramNoFrameskip-v4') # ddpg_continuous() - # ppo_continuous() + ppo_continuous() # action_conditional_video_prediction() diff --git a/network/base_network.py b/network/base_network.py index d3ffc23..9425a06 100644 --- a/network/base_network.py +++ b/network/base_network.py @@ -145,3 +145,44 @@ class TwoLayerFCNet(nn.Module): y = self.gate(self.fc1(x)) y = self.gate(self.fc2(y)) return y + +class ContinuousActorCriticWrapper: + def __init__(self, state_dim, action_dim, actor_fn, critic_fn, actor_opt_fn, critic_opt_fn): + self.actor = actor_fn(state_dim, action_dim) + self.critic = critic_fn(state_dim) + self.actor_opt = actor_opt_fn(self.actor.parameters()) + self.critic_opt = critic_opt_fn(self.critic.parameters()) + + def predict(self, state, actions=None): + mean, std, log_std = self.actor.predict(state) + values = self.critic.predict(state) + dist = torch.distributions.Normal(mean, std) + if actions is None: + actions = dist.sample() + log_probs = dist.log_prob(actions) + log_probs = torch.sum(log_probs, dim=1, keepdim=True) + return actions, log_probs, 0, values + + def variable(self, x, dtype=torch.FloatTensor): + return self.actor.variable(x, dtype) + + def tensor(self, x, dtype=torch.FloatTensor): + return self.actor.tensor(x, dtype) + + def zero_grad(self): + self.actor_opt.zero_grad() + self.critic_opt.zero_grad() + + def parameters(self): + return list(self.actor.parameters()) + list(self.critic.parameters()) + + def step(self): + self.actor_opt.step() + self.critic_opt.step() + + def state_dict(self): + return [self.actor.state_dict(), self.critic.state_dict()] + + def load_state_dict(self, state_dicts): + self.actor.load_state_dict(state_dicts[0]) + self.critic.load_state_dict(state_dicts[1])