I've written my own version of SAC(v2) for a problem with continuous action space. While training, the losses for the value network and both q functions steadily decrease down to 0.02-0.03. The loss for my actor/agent is negative and decreases to about -0.25 (I've read that it doesn't matter whether it is negative or not, but I'm not 100% sure). Despite that, the output variance from the Gaussian policy is way too high (make all the outcomes nearly uniformly likely) and is not decreasing during training. Does anyone know what can be a cause of that?
My implementation is mostly based on https://github.com/keiohta/tf2rl/blob/master/tf2rl/algos/sac.py, but I resigned from computing td_errors.
Here is the code (in case you need it).
import tensorflow as tf
from tensorflow.keras.layers import *
from src.anfis.anfis_layers import *
from src.model.sac_layer import *
from src.anfis.anfis_model import AnfisGD
hidden_activation = 'elu'
output_activation = 'linear'
class NetworkModel:
def __init__(self, training):
self.parameters_count = 2
self.results_count = 1
self.parameters_sets_count = [3, 4]
self.parameters_sets_total_count = sum(self.parameters_sets_count)
self.models = {}
self._initialise_layers() # initialises self.models[]
self.training = training
self.train()
def _initialise_layers(self):
# ------------
# LAYERS & DEBUG
# ------------
f_states = Input(shape=(self.parameters_count,))
f_actions = Input(shape=(self.results_count,))
# = tf.keras.layers.Dense(10)# AnfisGD(self.parameters_sets_count)
#f_anfis = model_anfis(densanf)#model_anfis(f_states)
f_policy_1 = tf.keras.layers.Dense(5, activation=hidden_activation)(f_states)
f_policy_2 = tf.keras.layers.Dense(5, activation=hidden_activation)(f_policy_1)
f_policy_musig = tf.keras.layers.Dense(2, activation=output_activation)(f_policy_2)
f_policy = GaussianLayer()(f_policy_musig)
#self.models["anfis"] = tf.keras.Model(inputs=f_states, outputs=f_anfis)
#self.models["forward"] = tf.keras.Model(inputs=f_states, outputs=model_anfis.anfis_forward(f_states))
self.models["actor"] = tf.keras.Model(inputs=f_states, outputs=f_policy)
self.models["critic-q-1"] = generate_q_network([f_states, f_actions])
self.models["critic-q-2"] = generate_q_network([f_states, f_actions])
self.models["critic-v"] = generate_value_network(f_states)
self.models["critic-v-t"] = generate_value_network(f_states)
# self.models["anfis"].compile(
# loss=tf.losses.mean_absolute_error,
# optimizer=tf.keras.optimizers.SGD(
# clipnorm=0.5,
# learning_rate=1e-3),
# metrics=[tf.keras.metrics.RootMeanSquaredError()]
# )
# self.models["forward"].compile(
# loss=tf.losses.mean_absolute_error,
# optimizer=tf.keras.optimizers.SGD(
# clipnorm=0.5,
# learning_rate=1e-3),
# metrics=[tf.keras.metrics.RootMeanSquaredError()]
# )
self.models["actor"].compile(
loss=tf.losses.mean_squared_error,
optimizer=tf.keras.optimizers.Adam(
learning_rate=1e-3),
metrics=[tf.keras.metrics.RootMeanSquaredError()]
)
def act(self, din):
data_input = tf.convert_to_tensor([din], dtype='float64')
data_output = self.models["actor"](data_input)[0]
return data_output.numpy()[0]
def train(self):
self.training.train(self, hybrid=False)
def mean(y_true, y_pred): #ignore y_pred
return tf.reduce_mean(y_true)
def generate_value_network(inputs):
# SAC Critic Value (Estimating rewards of being in state s)
f_critic_v1 = tf.keras.layers.Dense(5, activation=hidden_activation)(inputs)
f_critic_v2 = tf.keras.layers.Dense(5, activation=hidden_activation)(f_critic_v1)
f_critic_v = tf.keras.layers.Dense(1, activation=output_activation)(f_critic_v2)
m_value = tf.keras.Model(inputs=inputs, outputs=f_critic_v)
m_value.compile(
loss=tf.losses.mean_squared_error,
optimizer=tf.keras.optimizers.Adam(
learning_rate=1e-3),
metrics=[tf.keras.metrics.RootMeanSquaredError()]
)
return m_value
def generate_q_network(inputs):
# SAC Critic Q (Estimating rewards of taking action a while in state s)
f_critic_q_concatenate = tf.keras.layers.Concatenate()(inputs)
f_critic_q1 = tf.keras.layers.Dense(5, activation=hidden_activation)(f_critic_q_concatenate)
f_critic_q2 = tf.keras.layers.Dense(5, activation=hidden_activation)(f_critic_q1)
f_critic_q = tf.keras.layers.Dense(1, activation=output_activation)(f_critic_q2)
m_q = tf.keras.Model(inputs=inputs, outputs=f_critic_q)
m_q.compile(
loss=tf.losses.mean_squared_error,
optimizer=tf.keras.optimizers.Adam(
learning_rate=1e-3),
metrics=[tf.keras.metrics.RootMeanSquaredError()]
)
return m_q;
from src.model.training import Training
import numpy as np
import tensorflow as tf
from src.constructs.experience_holder import ExperienceHolder
class SACTraining(Training):
def __init__(self, environment):
super().__init__()
self.environment = environment
self.models = None
self.parameters_sets_count = None
self.parameters_sets_total_count = 0
self.parameters_count = 0
self.gamma = 0.99
self.alpha = 1.0
self.beta = 0.003
self.tau = 0.01
self.experience = ExperienceHolder(capacity=10000, cells=5) # state, action, reward, state', done
def train(self, simulation_model, **kwargs):
self.models = simulation_model.models
self.parameters_count = simulation_model.parameters_count
self.parameters_sets_count = simulation_model.parameters_sets_count
self.parameters_sets_total_count = simulation_model.parameters_sets_total_count
self.train_sac(
self.models,
epochs=300, max_steps=200, experience_batch=128, simulation=self.environment)
def train_sac(self, models, epochs, max_steps, experience_batch, simulation):
# deterministic random
np.random.seed(0)
history = []
epoch_steps = 128
simulation.reset()
update_net(models['critic-v'], models['critic-v-t'], 1.0)
for i in range(epochs):
print("epoch: ", i)
episode_reward = 0
reset = False
j = 0
while not(j > epoch_steps and reset):
j += 1
reset = False
# ---------------------------
# Observe state s and select action according to current policy
# ---------------------------
# Get simulation state
state_raw = simulation.get_normalised()
# state_unwound = [[i for t in state for i in t]]
state = [state_raw[0]] # TODO
state_tf = tf.convert_to_tensor(state)
# Get actions distribution from current model
# and their approx value from critic
actions_tf, _, _ = models['actor'](state_tf)
actions = list(actions_tf.numpy()[0])
# ---------------------------
# Execute action in the environment
# ---------------------------
reward, done = simulation.step_nominalised(actions)
episode_reward += reward
# ---------------------------
# Observe next state
# ---------------------------
state_l_raw = simulation.get_normalised()
state_l = [state_l_raw[0]] # TODO
# ---------------------------
# Store information in replay buffer
# ---------------------------
self.experience.save((state, actions, reward, state_l, 1 if not done else 0))
if done or simulation.step_counter > max_steps:
simulation.reset()
reset = True
# ---------------------------
# Updating network
# ---------------------------
if self.experience.size() > 500: # update_counter_limit:
exp = self.experience.replay(min(experience_batch, int(self.experience.size() * 0.8)))
states_tf = tf.convert_to_tensor(exp[0], dtype='float64')
actions_tf = tf.convert_to_tensor(exp[1], dtype='float64')
rewards_tf = tf.convert_to_tensor(exp[2], dtype='float64')
states_l_tf = tf.convert_to_tensor(exp[3], dtype='float64')
not_dones_tf = tf.convert_to_tensor(exp[4], dtype='float64')
with tf.GradientTape(watch_accessed_variables=True, persistent=True) as tape:
q_1_current = models['critic-q-1']([states_tf, actions_tf])
q_2_current = models['critic-q-2']([states_tf, actions_tf])
v_l_current = models['critic-v-t'](states_l_tf)
q_target = tf.stop_gradient(rewards_tf + not_dones_tf * self.gamma * v_l_current)
q_1_loss = tf.reduce_mean((q_target - q_1_current) ** 2)
q_2_loss = tf.reduce_mean((q_target - q_2_current) ** 2)
v_current = models['critic-v'](states_tf)
actions, policy_loss, sigma = models['actor'](states_tf)
q_1_policy = models['critic-q-1']([states_tf, actions_tf])
q_2_policy = models['critic-q-2']([states_tf, actions_tf])
q_min_policy = tf.minimum(q_1_policy, q_2_policy)
v_target = tf.stop_gradient(q_min_policy - self.alpha * policy_loss)
v_loss = tf.reduce_mean((v_target - v_current)**2)
a_loss = tf.reduce_mean(self.alpha * policy_loss - q_min_policy)
backward(tape, models['critic-q-1'], q_1_loss)
backward(tape, models['critic-q-2'], q_2_loss)
backward(tape, models['critic-v'], v_loss)
update_net(models['critic-v'], models['critic-v-t'], self.tau)
backward(tape, models['actor'], a_loss)
del tape
print('Loss:\n\tvalue: {}\n\tq1 : {}\n\tq2 : {}\n\tactor (ascent): {}'.format(
tf.reduce_mean(v_loss),
tf.reduce_mean(q_1_loss),
tf.reduce_mean(q_2_loss),
tf.reduce_mean(a_loss) #Gradient ascent
))
print('Episode Reward: {}'.format(episode_reward))
print('Batch sigma: {}'.format(tf.reduce_mean(sigma)))
def update_net(model, target, tau):
len_vars = len(model.trainable_variables)
for i in range(len_vars):
target.trainable_variables[i] = tau * model.trainable_variables[i] + (1.0 - tau) * target.trainable_variables[i]
def backward(tape, model, loss):
grads = tape.gradient(loss, model.trainable_variables)
model.optimizer.apply_gradients(
zip(grads, model.trainable_variables))
from tensorflow.keras import Model
import tensorflow as tf
import tensorflow_probability as tfp
class GaussianLayer(Model):
def __init__(self, **kwargs):
super(GaussianLayer, self).__init__(**kwargs)
def call(self, inputs, **kwargs):
mu, log_sig = tf.split(inputs, num_or_size_splits=2, axis=1)
log_sig_clip = tf.clip_by_value(log_sig, -20, 2)
sig = tf.exp(log_sig_clip)
distribution = tfp.distributions.Normal(mu, sig)
output = distribution.sample()
actions = tf.tanh(output)
return actions, \
distribution.log_prob(output) - \
tf.reduce_sum(tf.math.log(1 - actions ** 2 + 1e-12), axis=1, keepdims=True), \
tf.stop_gradient(tf.keras.backend.abs(actions - tf.tanh(mu)))