Tensorflow Reinforcement Learning RNN returning NaN's after Optimization with GradientTape

Question

def create_example_model():
    tf.keras.backend.set_floatx('float64')
    model = Sequential()
    model.add(LSTM(128, input_shape=((60, len(df_train.columns)))))

    model.add(Dense(64, activation='relu'))

    model.add(Dense(3, activation=None))

    return model

def choose_action(model, observation):
    observation = np.expand_dims(observation, axis=0)

    logits = model.predict(observation)

    prob_weights = tf.nn.softmax(logits).numpy()

    action = np.random.choice(3, size=1, p=prob_weights.flatten())[0]

    return action

def train_step(model, optimizer, observations, actions, discounted_rewards):
    with tf.GradientTape() as tape:

        logits = model(observations)

        loss = compute_loss(logits, actions, discounted_rewards)

        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

learning_rate = 1e-3
optimizer = tf.keras.optimizers.Adam(learning_rate)
env = TradingEnv(rnn_ready_array)

model = create_example_model()
memory = Memory()
info_list = []

for i_episode in range(10):
    observation = env.reset()
    memory.clear()

    while True:
        action = choose_action(model, observation)
        next_observation, reward, done, info = env.step(action)
        info_list.append(info)
        memory.add_to_memory(observation, action, reward)
        if done:
            total_reward = sum(memory.rewards)
            train_step(model, optimizer,
                 observations=np.array(memory.observations),
                 actions=np.array(memory.actions),
                 discounted_rewards = discount_rewards(memory.rewards))

            memory.clear()
            break
        observation = next_observation

I am working on a reinforcement learning project with Tensorflow 2.0; the format of the code comes from an online MIT course of which I am attempting to adapt to my own project. I am new to Tensorflow 2.0 and I can't glean from the documentation why this problem is occurring. The issue is that when I run the reinforcement learning process,

1. The first episode will always complete successfully.
2. A new observation will always be generated from the model successfully.
3. During the second episode, the network will always output: [NaN, NaN, NaN]

Some debugging info I have found that should be helpful: If I comment out the optimization lines 'grads = tape.gradient(...)' and 'optimizer.apply_gradients(...)' the script will run to completion error free (though it is obviously not doing anything useful without optimization). This indicates to me the optimization process is changing the model in a way that is causing the problem. I've tried to include only the necessary functions for debugging; if there is any further information one might need for debugging, I'd be happy to add additional info in an edit.

Answer 1

After hours of checking and rechecking various containers I realized it was the discounted rewards function that was not working properly, returning NaN in this circumstance. Issue resolved:)