lunar_lander_a2c_tdn_buffer_with_entropy.py

import gym
import numpy as np
import tensorflow as tf
from tensorflow import keras
from rl_utils.SARST_NStepReturn_RandomAccess_MemoryBuffer_NoOverwrite import SARST_NStepReturn_RandomAccess_MemoryBuffer_NoOverwrite

gpus = tf.config.list_physical_devices('GPU')
tf.config.set_visible_devices(gpus[0], 'GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)

env = gym.make('LunarLander-v2')

num_episodes = 10000
actor_learning_rate = 3e-4
critic_learning_rate = 3e-4
X_shape = (env.observation_space.shape[0])
gamma = 0.99
entropy_beta = 0.01
N_return = 3
batch_size = 16
log_std_min=-20
log_std_max=2

outputs_count = env.action_space.n

RND_SEED = 0x12345
tf.random.set_seed(RND_SEED)
np.random.random(RND_SEED)

rewards_history = []

actor_optimizer = tf.keras.optimizers.Adam(actor_learning_rate)
critic_optimizer = tf.keras.optimizers.Adam(critic_learning_rate)

exp_buffer = SARST_NStepReturn_RandomAccess_MemoryBuffer_NoOverwrite(1001, N_return, gamma, env.observation_space.shape, env.action_space.shape, np.int32)

def policy_network():
    input = keras.layers.Input(shape=(X_shape))
    x = keras.layers.Dense(256, activation='relu', kernel_initializer = keras.initializers.HeNormal(seed = RND_SEED),
                                                   bias_initializer = keras.initializers.HeNormal(seed = RND_SEED))(input)
    x = keras.layers.Dense(256, activation='relu', kernel_initializer = keras.initializers.HeNormal(seed = RND_SEED),
                                                   bias_initializer = keras.initializers.HeNormal(seed = RND_SEED))(x)
    actions_layer = keras.layers.Dense(outputs_count, activation='linear')(x)

    model = keras.Model(inputs=input, outputs=actions_layer)
    return model

def value_network():
    input = keras.layers.Input(shape=(X_shape))
    x = keras.layers.Dense(256, activation='relu', kernel_initializer = keras.initializers.HeNormal(seed = RND_SEED),
                                                   bias_initializer = keras.initializers.HeNormal(seed = RND_SEED))(input)
    x = keras.layers.Dense(256, activation='relu', kernel_initializer = keras.initializers.HeNormal(seed = RND_SEED),
                                                   bias_initializer = keras.initializers.HeNormal(seed = RND_SEED))(x)
    v_layer = keras.layers.Dense(1, activation='linear')(x)

    model = keras.Model(inputs=input, outputs=v_layer)
    return model

actor = policy_network()
actor_optimizer.build(actor.trainable_variables)
critic = value_network()
critic_optimizer.build(critic.trainable_variables)

@tf.function
def train_actor(states, actions, advantages):
    one_hot_actions_mask = tf.one_hot(actions, depth=outputs_count, on_value = 1.0, off_value = 0.0, dtype=tf.float32)
    with tf.GradientTape() as tape:
        actions_logits = actor(states, training=True)
        actions_logits = tf.clip_by_value(actions_logits, log_std_min, log_std_max)
        actions_log_distribution = tf.nn.log_softmax(actions_logits)

        actions_distribution = tf.nn.softmax(actions_logits)
        entropy = -tf.reduce_sum(actions_log_distribution * actions_distribution)

        loss = -tf.reduce_mean(tf.reduce_sum(actions_log_distribution * one_hot_actions_mask, axis=1) * advantages) + entropy_beta * entropy
    gradients = tape.gradient(loss, actor.trainable_variables)
    actor_optimizer.apply_gradients(zip(gradients, actor.trainable_variables))
    return loss

@tf.function
def train_critic(states, next_states, rewards, gamma_powers, dones):
    with tf.GradientTape() as tape:
        current_state_value = tf.squeeze(critic(states, training=True))
        next_state_values = tf.squeeze(critic(next_states, training=False))
        n_step_return = rewards + (1 - dones) * tf.math.pow(gamma, tf.cast(gamma_powers + 1, dtype=tf.float32)) * next_state_values
        
        advantage = n_step_return - current_state_value # TD(N) error
        loss = tf.reduce_mean(0.5 * tf.math.pow(advantage, 2)) #mse
    gradients = tape.gradient(loss, critic.trainable_variables)
    critic_optimizer.apply_gradients(zip(gradients, critic.trainable_variables))
    return loss, advantage

for i in range(num_episodes):
    observation, _ = env.reset()
    epoch_steps = 0
    done = False
    trunc = False
    episodic_reward = 0
    critic_loss_history = []
    actor_loss_history = []
    exp_buffer.reset()

    while not (done or trunc):
        actions_logits = actor(np.expand_dims(observation, axis = 0), training=False)[0]
        actions_distribution = tf.nn.softmax(actions_logits).numpy()
        chosen_action = np.random.choice(env.action_space.n, p=actions_distribution)
        
        next_observation, reward, done, trunc, _ = env.step(chosen_action)
        episodic_reward += reward
        
        exp_buffer.store(observation, 
                         chosen_action, 
                         next_observation, 
                         reward, 
                         float(done or trunc))
        epoch_steps+=1
        observation = next_observation

    learning_iterations = len(exp_buffer) // batch_size
    for l in range(learning_iterations + 1):
        if l == learning_iterations:
            states, actions, next_states, rewards, gammas, dones = exp_buffer.get_tail_batch(batch_size)
        else:
            states, actions, next_states, rewards, gammas, dones = exp_buffer(batch_size)
        critic_loss, adv = train_critic(states, next_states, rewards, gammas, dones)
        critic_loss_history.append(critic_loss)
        actor_loss = train_actor(states, actions, adv)
        actor_loss_history.append(actor_loss)

    rewards_history.append(episodic_reward)

    last_mean = np.mean(rewards_history[-100:])
    print(f'[epoch {i} ({epoch_steps})] Actor mloss: {np.mean(actor_loss_history):.4f} Critic mloss: {np.mean(critic_loss_history):.4f} Total reward: {episodic_reward} Mean(100)={last_mean:.4f}')
    if last_mean > 200:
        break
env.close()
if last_mean > 200:
    actor.save('lunar_lander_a2c_nrH.h5')
input("training complete...")