Based on this implementation of TD3, I converted different torch-based methods to their tensorflow equivalents. No matter what hyperparameters I use, torch's version will start converging by 50,000 - 100,000 steps which happens within a few minutes of starting the training. However, tensorflow version doesn't show any signs of improvement up to 600,000 steps. Here's a notebook with both versions and I will include here the converted methods. I'm using gym's BipedalWalker-v3 environment for both versions.
What I tried
I tried everything from lowering the default learning rate of the paper 0.0003, hidden unit sizes (i tried 64, 128, 256, 512 and even 400, 300 the paper's defaults nothing works), changed batch size (100, 128, 256, ...), lowered buffer size (default: 1,000,000) to 100,000, 50,000 ... I also tried reusing tf.GradientTape(persistent=True) and doing the updates using the same tape, tried using tf.train.ExponentialMovingAverage instead of manual averaging and needless to say, I even re-wrote the code multiple times. NOTHING works.
Note
self.max_action in torch's implementation, equals 1 and you will find it hardcoded in tf's, to clear any confusion you may have.
Common hyperparameters
- dense units: 256 (different from paper's 400 and 300)
- actor update delay: 2
- exploration noise std: 0.1
- policy noise std: 0.2
- noise clip norm: 0.5
- tau: 0.005
- optimizer: adam
- learning rate: 0.0003
- replay buffer size: 1,000,000
- replay initial size: 25,000
- batch size: 256 (different from paper's 100)
Torch actor and critic models
import torch
import torch.nn as nn
import torch.nn.functional as F
class Actor(nn.Module):
def __init__(self, state_dim, action_dim, max_action):
super(Actor, self).__init__()
self.l1 = nn.Linear(state_dim, 256)
self.l2 = nn.Linear(256, 256)
self.l3 = nn.Linear(256, action_dim)
self.max_action = max_action
def forward(self, state):
a = F.relu(self.l1(state))
a = F.relu(self.l2(a))
return self.max_action * torch.tanh(self.l3(a))
class Critic(nn.Module):
def __init__(self, state_dim, action_dim):
super(Critic, self).__init__()
self.l1 = nn.Linear(state_dim + action_dim, 256)
self.l2 = nn.Linear(256, 256)
self.l3 = nn.Linear(256, 1)
self.l4 = nn.Linear(state_dim + action_dim, 256)
self.l5 = nn.Linear(256, 256)
self.l6 = nn.Linear(256, 1)
def forward(self, state, action):
sa = torch.cat([state, action], 1)
q1 = F.relu(self.l1(sa))
q1 = F.relu(self.l2(q1))
q1 = self.l3(q1)
q2 = F.relu(self.l4(sa))
q2 = F.relu(self.l5(q2))
q2 = self.l6(q2)
return q1, q2
def Q1(self, state, action):
sa = torch.cat([state, action], 1)
q1 = F.relu(self.l1(sa))
q1 = F.relu(self.l2(q1))
q1 = self.l3(q1)
return q1
Tensorflow actor and critic models
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Input
def actor(input_shape, output_units):
x0 = Input(input_shape)
x = Dense(256, 'relu')(x0)
x = Dense(256, 'relu')(x)
output = Dense(output_units, 'tanh')(x)
return Model(x0, output)
def critic(input_shape):
x0 = Input(input_shape)
x1 = Dense(256, 'relu')(x0)
x1 = Dense(256, 'relu')(x1)
v1 = Dense(1)(x1)
x2 = Dense(256, 'relu')(x0)
x2 = Dense(256, 'relu')(x2)
v2 = Dense(1)(x2)
return Model(x0, [v1, v2])
Torch train step
def train(self, replay_buffer, batch_size=256):
self.total_it += 1
state, action, next_state, reward, not_done = replay_buffer.sample(batch_size)
with torch.no_grad():
noise = (torch.randn_like(action) * self.policy_noise).clamp(
-self.noise_clip, self.noise_clip
)
next_action = (self.actor_target(next_state) + noise).clamp(
-self.max_action, self.max_action
)
target_Q1, target_Q2 = self.critic_target(next_state, next_action)
target_Q = torch.min(target_Q1, target_Q2)
target_Q = reward + not_done * self.discount * target_Q
current_Q1, current_Q2 = self.critic(state, action)
critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(
current_Q2, target_Q
)
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
if self.total_it % self.policy_freq == 0:
actor_loss = -self.critic.Q1(state, self.actor(state)).mean()
self.actor_optimizer.zero_grad()
actor_loss.backward()
self.actor_optimizer.step()
for param, target_param in zip(
self.critic.parameters(), self.critic_target.parameters()
):
target_param.data.copy_(
self.tau * param.data + (1 - self.tau) * target_param.data
)
for param, target_param in zip(
self.actor.parameters(), self.actor_target.parameters()
):
target_param.data.copy_(
self.tau * param.data + (1 - self.tau) * target_param.data
)
tensorflow training methods
def update_critic_weights(self, states, actions, rewards, dones, new_states):
policy_noise = tf.random.normal(
(self.n_envs * self.buffer_batch_size, self.n_actions),
0,
self.policy_noise_std,
)
policy_noise = tf.clip_by_value(policy_noise, -self.noise_clip, self.noise_clip)
target_actions = tf.clip_by_value(
self.target_actor(new_states) + policy_noise, -1, 1
)
target_value1, target_value2 = self.target_critic(
tf.concat([new_states, target_actions], 1)
)
target_value = rewards + (1 - dones) * self.gamma * tf.minimum(
target_value1, target_value2
)
with tf.GradientTape() as tape:
value1, value2 = self.critic(tf.concat([states, actions], 1))
value_loss = MSE(value1, target_value) + MSE(value2, target_value2)
self.critic.optimizer.minimize(
value_loss, self.critic.trainable_variables, tape=tape
)
def update_actor_weights(self, states):
with tf.GradientTape() as tape:
actions = self.actor(states)
loss = -tf.reduce_mean(self.critic([tf.concat([states, actions], 1)])[0])
self.actor.optimizer.minimize(loss, self.actor.trainable_variables, tape=tape)
@tf.function
def train_step(self):
self.updates.assign_add(1)
step_actions = self.actor(self.get_states())
exploration_noise = tf.random.normal(
[self.n_actions], 0, self.exploration_noise_std
)
step_actions = tf.clip_by_value(step_actions + exploration_noise, -1, 1)
tf.numpy_function(self.step_envs, [step_actions, False, True], [])
batch = tf.numpy_function(self.concat_buffer_samples, [], 5 * [tf.float32])
self.update_critic_weights(*batch)
if tf.math.equal(tf.math.mod(self.updates, self.actor_delay), 0):
self.update_actor_weights(batch[0])
self.sync_target_models()
def sync_target_models(self):
for model, target_model in self.model_groups:
for var, target_var in zip(
model.trainable_variables, target_model.trainable_variables
):
target_var.assign(self.tau * var + (1 - self.tau) * target_var)
The results which you can reproduce using the notebook:
Torch
time: 0:00:01.928306, steps: 3875, games: 10, fps: 2009.0, mean reward: -105.27, best reward: -82.24
time: 0:00:06.234093, steps: 13832, games: 20, fps: 2312.0, mean reward: -99.73, best reward: -77.94
time: 0:00:09.219966, steps: 20691, games: 30, fps: 2297.0, mean reward: -99.37, best reward: -75.26
time: 0:00:24.854520, steps: 27736, games: 40, fps: 450.0, mean reward: -102.14, best reward: -65.01
time: 0:00:28.400598, steps: 28444, games: 50, fps: 199.0, mean reward: -103.27, best reward: -65.01
time: 0:01:24.811760, steps: 39761, games: 60, fps: 200.0, mean reward: -107.92, best reward: -65.01
time: 0:02:23.950100, steps: 51719, games: 70, fps: 202.0, mean reward: -105.17, best reward: -65.01
time: 0:03:29.150313, steps: 64853, games: 80, fps: 201.0, mean reward: -103.09, best reward: -65.01
time: 0:04:27.118459, steps: 76666, games: 90, fps: 203.0, mean reward: -105.1, best reward: -65.01
time: 0:05:33.525320, steps: 90105, games: 100, fps: 202.0, mean reward: -104.47, best reward: -31.01
time: 0:06:38.586421, steps: 103162, games: 110, fps: 200.0, mean reward: -102.21, best reward: -31.01
time: 0:07:50.672645, steps: 117707, games: 120, fps: 201.0, mean reward: -99.69, best reward: -31.01
time: 0:09:10.076208, steps: 133707, games: 130, fps: 201.0, mean reward: -93.53, best reward: -31.01
time: 0:10:07.278960, steps: 145180, games: 140, fps: 200.0, mean reward: -88.81, best reward: -30.43
time: 0:11:04.842319, steps: 156767, games: 150, fps: 201.0, mean reward: -85.12, best reward: -30.43
time: 0:11:30.127616, steps: 161788, games: 160, fps: 198.0, mean reward: -83.37, best reward: -30.43
time: 0:11:34.950709, steps: 162755, games: 170, fps: 200.0, mean reward: -86.65, best reward: -30.43
time: 0:12:06.615783, steps: 169146, games: 180, fps: 201.0, mean reward: -89.97, best reward: -30.43
time: 0:13:11.028382, steps: 182064, games: 190, fps: 200.0, mean reward: -84.38, best reward: -30.43
time: 0:14:30.357974, steps: 198064, games: 200, fps: 201.0, mean reward: -79.24, best reward: -30.43
time: 0:15:49.531968, steps: 214064, games: 210, fps: 202.0, mean reward: -74.81, best reward: -27.15
time: 0:17:01.859939, steps: 228590, games: 220, fps: 200.0, mean reward: -71.18, best reward: 37.19
Tensorflow:
time: 0:00:33.882932, steps: 6825, games: 10, fps: 206, mean reward: -81.9, best reward: -68.04
time: 0:01:30.612159, steps: 18306, games: 20, fps: 201, mean reward: -77.51, best reward: -68.04
time: 0:01:49.359029, steps: 22075, games: 30, fps: 190, mean reward: -88.67, best reward: -68.04
time: 0:01:53.394437, steps: 22881, games: 40, fps: 190, mean reward: -94.84, best reward: -68.04
time: 0:01:56.996669, steps: 23603, games: 50, fps: 205, mean reward: -98.51, best reward: -68.04
time: 0:02:16.071768, steps: 27402, games: 60, fps: 192, mean reward: -97.91, best reward: -68.04
time: 0:02:28.969803, steps: 29996, games: 70, fps: 204, mean reward: -100.71, best reward: -68.04
time: 0:02:32.286386, steps: 30656, games: 80, fps: 203, mean reward: -102.14, best reward: -68.04
time: 0:02:41.094522, steps: 32407, games: 90, fps: 198, mean reward: -103.45, best reward: -68.04
time: 0:02:57.489013, steps: 35634, games: 100, fps: 171, mean reward: -104.67, best reward: -68.04
time: 0:03:15.559559, steps: 39184, games: 110, fps: 199, mean reward: -109.12, best reward: -68.04
time: 0:03:26.770144, steps: 41374, games: 120, fps: 194, mean reward: -113.48, best reward: -68.04
time: 0:03:56.878539, steps: 47205, games: 130, fps: 195, mean reward: -114.48, best reward: -68.04
time: 0:04:37.541307, steps: 55078, games: 140, fps: 194, mean reward: -116.36, best reward: -68.04
time: 0:05:14.164399, steps: 62089, games: 150, fps: 192, mean reward: -115.31, best reward: -68.04
time: 0:06:14.816104, steps: 73559, games: 160, fps: 197, mean reward: -118.41, best reward: -68.04
time: 0:06:51.555758, steps: 80434, games: 170, fps: 185, mean reward: -118.5, best reward: -68.04
time: 0:07:11.885412, steps: 84218, games: 180, fps: 188, mean reward: -118.14, best reward: -68.04
time: 0:07:51.662705, steps: 91584, games: 190, fps: 184, mean reward: -119.91, best reward: -68.04
time: 0:08:12.492280, steps: 95425, games: 200, fps: 181, mean reward: -119.18, best reward: -68.04
time: 0:08:33.173622, steps: 99274, games: 210, fps: 186, mean reward: -117.43, best reward: -68.04
time: 0:08:40.747114, steps: 100642, games: 220, fps: 189, mean reward: -116.43, best reward: -68.04
time: 0:09:09.865704, steps: 105968, games: 230, fps: 185, mean reward: -116.34, best reward: -68.04
And it will keep diverging forever guaranteed, I almost reached 1,000,000 steps and the reward keeps circling around the same average.