From f4fe1362f05edd0ec1109de5e0b4906e3ee55841 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 24 Sep 2019 14:53:03 +0200 Subject: [PATCH] Renaming --- torchy_baselines/common/buffers.py | 44 +++++++++++++++--------------- torchy_baselines/ppo/policies.py | 33 +++++++++++----------- torchy_baselines/ppo/ppo.py | 12 ++++---- torchy_baselines/sac/policies.py | 26 +++++++++--------- torchy_baselines/sac/sac.py | 16 +++++------ torchy_baselines/td3/policies.py | 21 ++++++++------ torchy_baselines/td3/td3.py | 21 +++++++------- 7 files changed, 88 insertions(+), 85 deletions(-) diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 599338a..61421fa 100644 --- a/torchy_baselines/common/buffers.py +++ b/torchy_baselines/common/buffers.py @@ -3,10 +3,10 @@ import torch as th class BaseBuffer(object): - def __init__(self, buffer_size, state_dim, action_dim, device='cpu', n_envs=1): + def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1): super(BaseBuffer, self).__init__() self.buffer_size = buffer_size - self.state_dim = state_dim + self.obs_dim = obs_dim self.action_dim = action_dim self.pos = 0 self.full = False @@ -58,20 +58,20 @@ class ReplayBuffer(BaseBuffer): Taken from https://github.com/apourchot/CEM-RL """ - def __init__(self, buffer_size, state_dim, action_dim, device='cpu', n_envs=1): - super(ReplayBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs) + def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1): + super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs) assert n_envs == 1 - self.states = th.zeros(self.buffer_size, self.n_envs, self.state_dim) + self.observations = th.zeros(self.buffer_size, self.n_envs, self.obs_dim) self.actions = th.zeros(self.buffer_size, self.n_envs, self.action_dim) - self.next_states = th.zeros(self.buffer_size, self.n_envs, self.state_dim) + self.next_observations = th.zeros(self.buffer_size, self.n_envs, self.obs_dim) self.rewards = th.zeros(self.buffer_size, self.n_envs) self.dones = th.zeros(self.buffer_size, self.n_envs) - def add(self, state, next_state, action, reward, done): + def add(self, obs, next_obs, action, reward, done): # Copy to avoid modification by reference - self.states[self.pos] = th.FloatTensor(np.array(state).copy()) - self.next_states[self.pos] = th.FloatTensor(np.array(next_state).copy()) + self.observations[self.pos] = th.FloatTensor(np.array(obs).copy()) + self.next_observations[self.pos] = th.FloatTensor(np.array(next_obs).copy()) self.actions[self.pos] = th.FloatTensor(np.array(action).copy()) self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy()) self.dones[self.pos] = th.FloatTensor(np.array(done).copy()) @@ -82,27 +82,27 @@ class ReplayBuffer(BaseBuffer): self.pos = 0 def _get_samples(self, batch_inds): - return (self.states[batch_inds, 0, :].to(self.device), + return (self.observations[batch_inds, 0, :].to(self.device), self.actions[batch_inds, 0, :].to(self.device), - self.next_states[batch_inds, 0, :].to(self.device), + self.next_observations[batch_inds, 0, :].to(self.device), self.dones[batch_inds].to(self.device), self.rewards[batch_inds].to(self.device)) class RolloutBuffer(BaseBuffer): - def __init__(self, buffer_size, state_dim, action_dim, device='cpu', - lambda_=1, gamma=0.99, n_envs=1): - super(RolloutBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs) + def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', + gae_lambda=1, gamma=0.99, n_envs=1): + super(RolloutBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs) # TODO: try the buffer on the gpu? - self.lambda_ = lambda_ + self.gae_lambda = gae_lambda self.gamma = gamma - self.states, self.actions, self.rewards, self.advantages = None, None, None, None + self.observations, self.actions, self.rewards, self.advantages = None, None, None, None self.returns, self.dones, self.values, self.log_probs = None, None, None, None self.generator_ready = False self.reset() def reset(self): - self.states = th.zeros(self.buffer_size, self.n_envs, self.state_dim) + self.observations = th.zeros(self.buffer_size, self.n_envs, self.obs_dim) self.actions = th.zeros(self.buffer_size, self.n_envs, self.action_dim) self.rewards = th.zeros(self.buffer_size, self.n_envs) self.returns = th.zeros(self.buffer_size, self.n_envs) @@ -126,12 +126,12 @@ class RolloutBuffer(BaseBuffer): next_non_terminal = 1.0 - self.dones[step + 1] next_value = self.values[step + 1] delta = self.rewards[step] + self.gamma * next_value * next_non_terminal - self.values[step] - last_gae_lam = delta + self.gamma * self.lambda_ * next_non_terminal * last_gae_lam + last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam self.advantages[step] = last_gae_lam self.returns = self.advantages + self.values - def add(self, state, action, reward, done, value, log_prob): - self.states[self.pos] = th.FloatTensor(np.array(state).copy()) + def add(self, obs, action, reward, done, value, log_prob): + self.observations[self.pos] = th.FloatTensor(np.array(obs).copy()) self.actions[self.pos] = th.FloatTensor(np.array(action).copy()) self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy()) self.dones[self.pos] = th.FloatTensor(np.array(done).copy()) @@ -146,7 +146,7 @@ class RolloutBuffer(BaseBuffer): indices = th.randperm(self.buffer_size * self.n_envs) # Prepare the data if not self.generator_ready: - for tensor in ['states', 'actions', 'values', + for tensor in ['observations', 'actions', 'values', 'log_probs', 'advantages', 'returns']: self.__dict__[tensor] = self.swap_and_flatten(self.__dict__[tensor]) self.generator_ready = True @@ -157,7 +157,7 @@ class RolloutBuffer(BaseBuffer): start_idx += batch_size def _get_samples(self, batch_inds): - return (self.states[batch_inds].to(self.device), + return (self.observations[batch_inds].to(self.device), self.actions[batch_inds].to(self.device), self.values[batch_inds].flatten().to(self.device), self.log_probs[batch_inds].flatten().to(self.device), diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index bc0d48b..e773251 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -2,18 +2,18 @@ from functools import partial import torch as th import torch.nn as nn -from torch.distributions import Normal import numpy as np from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution + class PPOPolicy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate=1e-3, net_arch=None, device='cpu', activation_fn=nn.Tanh, adam_epsilon=1e-5): super(PPOPolicy, self).__init__(observation_space, action_space, device) - self.state_dim = self.observation_space.shape[0] + self.obs_dim = self.observation_space.shape[0] self.action_dim = self.action_space.shape[0] if net_arch is None: net_arch = [64, 64] @@ -21,7 +21,7 @@ class PPOPolicy(BasePolicy): self.activation_fn = activation_fn self.adam_epsilon = adam_epsilon self.net_args = { - 'input_dim': self.state_dim, + 'input_dim': self.obs_dim, 'output_dim': -1, 'net_arch': self.net_arch, 'activation_fn': self.activation_fn @@ -41,12 +41,12 @@ class PPOPolicy(BasePolicy): def _build(self, learning_rate): # TODO: support shared network - # shared_net = create_mlp(self.state_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn) + # shared_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn) # self.shared_net = nn.Sequential(*shared_net).to(self.device) - pi_net = create_mlp(self.state_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn) + pi_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn) self.pi_net = nn.Sequential(*pi_net).to(self.device) - vf_net = create_mlp(self.state_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn) + vf_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn) self.vf_net = nn.Sequential(*vf_net).to(self.device) # self.action_net = nn.Linear(self.net_arch[-1], self.action_dim) @@ -67,32 +67,33 @@ class PPOPolicy(BasePolicy): # TODO: support linear decay of the learning rate self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate, eps=self.adam_epsilon) - def forward(self, state, deterministic=False): - state = th.FloatTensor(state).to(self.device) - latent_pi, latent_vf = self._get_latent(state) + def forward(self, obs, deterministic=False): + if not isinstance(obs, th.Tensor): + obs = th.FloatTensor(obs).to(self.device) + latent_pi, latent_vf = self._get_latent(obs) value = self.value_net(latent_vf) action, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) log_prob = action_distribution.log_prob(action) return action, value, log_prob - def _get_latent(self, state): + def _get_latent(self, obs): if self.shared_net is not None: - latent = self.shared_net(state) + latent = self.shared_net(obs) return latent, latent else: - return self.pi_net(state), self.vf_net(state) + return self.pi_net(obs), self.vf_net(obs) def _get_action_dist_from_latent(self, latent, deterministic=False): mean_actions = self.action_net(latent) return self.action_dist.proba_distribution(mean_actions, self.log_std, deterministic=deterministic) - def actor_forward(self, state, deterministic=False): - latent_pi, _ = self._get_latent(state) + def actor_forward(self, obs, deterministic=False): + latent_pi, _ = self._get_latent(obs) action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) return action.detach().cpu().numpy() - def get_policy_stats(self, state, action): - latent_pi, latent_vf = self._get_latent(state) + def get_policy_stats(self, obs, action): + latent_pi, latent_vf = self._get_latent(obs) _, action_distribution = self._get_action_dist_from_latent(latent_pi) log_prob = action_distribution.log_prob(action) value = self.value_net(latent_vf) diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 8d5d902..bcb957c 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -31,7 +31,7 @@ class PPO(BaseRLModel): def __init__(self, policy, env, policy_kwargs=None, verbose=0, learning_rate=3e-4, seed=0, device='auto', n_optim=5, batch_size=64, n_steps=256, - gamma=0.99, lambda_=0.95, clip_range=0.2, + gamma=0.99, gae_lambda=0.95, clip_range=0.2, ent_coef=0.01, vf_coef=0.5, max_grad_norm=0.5, target_kl=None, clip_range_vf=None, create_eval_env=False, tensorboard_log=None, @@ -46,7 +46,7 @@ class PPO(BaseRLModel): self.n_optim = n_optim self.n_steps = n_steps self.gamma = gamma - self.lambda_ = lambda_ + self.gae_lambda = gae_lambda self.clip_range = clip_range self.ent_coef = ent_coef self.vf_coef = vf_coef @@ -67,7 +67,7 @@ class PPO(BaseRLModel): self.seed(self._seed) self.rollout_buffer = RolloutBuffer(self.n_steps, state_dim, action_dim, self.device, - gamma=self.gamma, lambda_=self.lambda_, n_envs=self.n_envs) + gamma=self.gamma, gae_lambda=self.gae_lambda, n_envs=self.n_envs) self.policy = self.policy(self.observation_space, self.action_space, self.learning_rate, device=self.device, **self.policy_kwargs) self.policy = self.policy.to(self.device) @@ -125,8 +125,8 @@ class PPO(BaseRLModel): # Sample replay buffer for replay_data in self.rollout_buffer.get(batch_size): # Unpack - state, action, old_values, old_log_prob, advantage, return_batch = replay_data - values, log_prob, entropy = self.policy.get_policy_stats(state, action) + obs, action, old_values, old_log_prob, advantage, return_batch = replay_data + values, log_prob, entropy = self.policy.get_policy_stats(obs, action) values = values.flatten() # Normalize advantage advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8) @@ -145,7 +145,7 @@ class PPO(BaseRLModel): # Clip the different between old and new value # NOTE: this depends on the reward scaling values_pred = old_values + th.clamp(values - old_values, -self.clip_range_vf, self.clip_range_vf) - # Value loss using the TD(lambda_) target + # Value loss using the TD(gae_lambda) target value_loss = F.mse_loss(return_batch, values_pred) diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index c89004e..34a1534 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -10,51 +10,51 @@ LOG_STD_MIN = -20 class Actor(BaseNetwork): - def __init__(self, state_dim, action_dim, net_arch=None, activation_fn=nn.ReLU): + def __init__(self, obs_dim, action_dim, net_arch=None, activation_fn=nn.ReLU): super(Actor, self).__init__() if net_arch is None: net_arch = [256, 256] # TODO: orthogonal initialization? - actor_net = create_mlp(state_dim, -1, net_arch, activation_fn) + actor_net = create_mlp(obs_dim, -1, net_arch, activation_fn) self.actor_net = nn.Sequential(*actor_net) self.action_dist = SquashedDiagGaussianDistribution(action_dim) self.mu = nn.Linear(net_arch[-1], action_dim) self.log_std = nn.Linear(net_arch[-1], action_dim) - def get_action_dist_params(self, state): - latent = self.actor_net(state) + def get_action_dist_params(self, obs): + latent = self.actor_net(obs) mean_actions, log_std = self.mu(latent), self.log_std(latent) # Original Implementation to cap the standard deviation log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX) return mean_actions, log_std - def forward(self, state, deterministic=False): - mean_actions, log_std = self.get_action_dist_params(state) + def forward(self, obs, deterministic=False): + mean_actions, log_std = self.get_action_dist_params(obs) # Note the action is squashed action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic) return action - def action_log_prob(self, state): - mean_actions, log_std = self.get_action_dist_params(state) + def action_log_prob(self, obs): + mean_actions, log_std = self.get_action_dist_params(obs) action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std) return action, log_prob class Critic(BaseNetwork): - def __init__(self, state_dim, action_dim, + def __init__(self, obs_dim, action_dim, net_arch=None, activation_fn=nn.ReLU): super(Critic, self).__init__() if net_arch is None: net_arch = [256, 256] - q1_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn) + q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) self.q1_net = nn.Sequential(*q1_net) - q2_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn) + q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) self.q2_net = nn.Sequential(*q2_net) self.q_networks = [self.q1_net, self.q2_net] @@ -72,12 +72,12 @@ class SACPolicy(BasePolicy): learning_rate=1e-3, net_arch=None, device='cpu', activation_fn=nn.ReLU): super(SACPolicy, self).__init__(observation_space, action_space, device) - self.state_dim = self.observation_space.shape[0] + self.obs_dim = self.observation_space.shape[0] self.action_dim = self.action_space.shape[0] self.net_arch = net_arch self.activation_fn = activation_fn self.net_args = { - 'state_dim': self.state_dim, + 'obs_dim': self.obs_dim, 'action_dim': self.action_dim, 'net_arch': self.net_arch, 'activation_fn': self.activation_fn diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 01a9d09..dc98937 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -52,7 +52,7 @@ class SAC(BaseRLModel): self._setup_model() def _setup_model(self): - state_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] + obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] self.seed(self._seed) # Target entropy is used when learning the entropy coefficient @@ -88,7 +88,7 @@ class SAC(BaseRLModel): # is passed self.ent_coef = float(self.ent_coef) - self.replay_buffer = ReplayBuffer(self.buffer_size, state_dim, action_dim, self.device) + self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device) self.policy = self.policy(self.observation_space, self.action_space, self.learning_rate, device=self.device, **self.policy_kwargs) self.policy = self.policy.to(self.device) @@ -125,10 +125,10 @@ class SAC(BaseRLModel): # Sample replay buffer replay_data = self.replay_buffer.sample(batch_size) - state, action_batch, next_state, done, reward = replay_data + obs, action_batch, next_obs, done, reward = replay_data # Action by the current actor for the sampled state - action_pi, log_prob = self.actor.action_log_prob(state) + action_pi, log_prob = self.actor.action_log_prob(obs) log_prob = log_prob.reshape(-1, 1) ent_coef_loss = None @@ -143,10 +143,10 @@ class SAC(BaseRLModel): self.ent_coef_optimizer.step() # Select action according to policy - next_action, next_log_prob = self.actor.action_log_prob(next_state) + next_action, next_log_prob = self.actor.action_log_prob(next_obs) # Compute the target Q value - target_q1, target_q2 = self.critic_target(next_state, next_action) + target_q1, target_q2 = self.critic_target(next_obs, next_action) target_q = th.min(target_q1, target_q2) target_q = reward + ((1 - done) * self.gamma * target_q).detach() @@ -155,7 +155,7 @@ class SAC(BaseRLModel): # Get current Q estimates # using action from the replay buffer - current_q1, current_q2 = self.critic(state, action_batch) + current_q1, current_q2 = self.critic(obs, action_batch) # Compute critic loss critic_loss = 0.5 * (F.mse_loss(current_q1, q_backup) + F.mse_loss(current_q2, q_backup)) @@ -167,7 +167,7 @@ class SAC(BaseRLModel): # Compute actor loss # Alternative: actor_loss = th.mean(log_prob - min_qf_pi) - actor_loss = (self.ent_coef * log_prob - self.critic.q1_forward(state, action_pi)).mean() + actor_loss = (self.ent_coef * log_prob - self.critic.q1_forward(obs, action_pi)).mean() # Optimize the actor self.actor.optimizer.zero_grad() diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 41a4371..5518b83 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -5,32 +5,32 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create class Actor(BaseNetwork): - def __init__(self, state_dim, action_dim, net_arch=None, activation_fn=nn.ReLU): + def __init__(self, obs_dim, action_dim, net_arch=None, activation_fn=nn.ReLU): super(Actor, self).__init__() if net_arch is None: net_arch = [400, 300] # TODO: orthogonal initialization? - actor_net = create_mlp(state_dim, action_dim, net_arch, activation_fn, squash_out=True) + actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) self.actor_net = nn.Sequential(*actor_net) - def forward(self, x): - return self.actor_net(x) + def forward(self, obs): + return self.actor_net(obs) class Critic(BaseNetwork): - def __init__(self, state_dim, action_dim, + def __init__(self, obs_dim, action_dim, net_arch=None, activation_fn=nn.ReLU): super(Critic, self).__init__() if net_arch is None: net_arch = [400, 300] - q1_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn) + q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) self.q1_net = nn.Sequential(*q1_net) - q2_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn) + q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) self.q2_net = nn.Sequential(*q2_net) self.q_networks = [self.q1_net, self.q2_net] @@ -48,12 +48,12 @@ class TD3Policy(BasePolicy): learning_rate=1e-3, net_arch=None, device='cpu', activation_fn=nn.ReLU): super(TD3Policy, self).__init__(observation_space, action_space, device) - self.state_dim = self.observation_space.shape[0] + self.obs_dim = self.observation_space.shape[0] self.action_dim = self.action_space.shape[0] self.net_arch = net_arch self.activation_fn = activation_fn self.net_args = { - 'state_dim': self.state_dim, + 'obs_dim': self.obs_dim, 'action_dim': self.action_dim, 'net_arch': self.net_arch, 'activation_fn': self.activation_fn @@ -79,6 +79,9 @@ class TD3Policy(BasePolicy): def make_critic(self): return Critic(**self.net_args).to(self.device) + def forward(self, obs): + return self.actor(obs) + MlpPolicy = TD3Policy diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 76aa9f1..2733147 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -16,7 +16,6 @@ class TD3(BaseRLModel): Paper: https://arxiv.org/abs/1802.09477 Code: https://github.com/sfujim/TD3 """ - def __init__(self, policy, env, policy_kwargs=None, verbose=0, buffer_size=int(1e6), learning_rate=1e-3, seed=0, device='auto', action_noise_std=0.1, start_timesteps=100, policy_freq=2, @@ -39,9 +38,9 @@ class TD3(BaseRLModel): self._setup_model() def _setup_model(self): - state_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] + obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] self.seed(self._seed) - self.replay_buffer = ReplayBuffer(self.buffer_size, state_dim, action_dim, self.device) + self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device) self.policy = self.policy(self.observation_space, self.action_space, self.learning_rate, device=self.device, **self.policy_kwargs) self.policy = self.policy.to(self.device) @@ -78,22 +77,22 @@ class TD3(BaseRLModel): for it in range(n_iterations): # Sample replay buffer if replay_data is None: - state, action, next_state, done, reward = self.replay_buffer.sample(batch_size) + obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size) else: - state, action, next_state, done, reward = replay_data + obs, action, next_obs, done, reward = replay_data # Select action according to policy and add clipped noise noise = action.clone().data.normal_(0, policy_noise) noise = noise.clamp(-noise_clip, noise_clip) - next_action = (self.actor_target(next_state) + noise).clamp(-1, 1) + next_action = (self.actor_target(next_obs) + noise).clamp(-1, 1) # Compute the target Q value - target_q1, target_q2 = self.critic_target(next_state, next_action) + target_q1, target_q2 = self.critic_target(next_obs, next_action) target_q = th.min(target_q1, target_q2) target_q = reward + ((1 - done) * discount * target_q).detach() # Get current Q estimates - current_q1, current_q2 = self.critic(state, action) + current_q1, current_q2 = self.critic(obs, action) # Compute critic loss critic_loss = F.mse_loss(current_q1, target_q) + F.mse_loss(current_q2, target_q) @@ -115,12 +114,12 @@ class TD3(BaseRLModel): for it in range(n_iterations): # Sample replay buffer if replay_data is None: - state, _, next_state, done, reward = self.replay_buffer.sample(batch_size) + obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size) else: - state, _, next_state, done, reward = replay_data + obs, _, next_obs, done, reward = replay_data # Compute actor loss - actor_loss = -self.critic.q1_forward(state, self.actor(state)).mean() + actor_loss = -self.critic.q1_forward(obs, self.actor(obs)).mean() # Optimize the actor self.actor.optimizer.zero_grad()