Variance of the Gaussian policy is not decreasing while training the agent using Soft Actor-Critic method

Question

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)))