check error in the given program

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

class PPO:
    def __init__(self, env, actor_lr=1e-4, critic_lr=1e-3, gamma=0.99, lam=0.95, clip_range=0.2):
        self.env = env
        self.actor_lr = actor_lr
        self.critic_lr = critic_lr
        self.gamma = gamma
        self.lam = lam
        self.clip_range = clip_range

        self.actor = Actor(env.observation_space, env.action_space)
        self.critic = Critic(env.observation_space)

        self.actor_optimizer = tf.keras.optimizers.Adam(learning_rate=actor_lr)
        self.critic_optimizer = tf.keras.optimizers.Adam(learning_rate=critic_lr)

    def collect_trajectories(self, episodes=10):
        trajectories = []
        for _ in range(episodes):
            trajectory = []
            state = self.env.reset()
            done = False
            while not done:
                action, _ = self.actor(np.expand_dims(state, axis=0))
                action = action.numpy()[0]
                next_state, reward, done, _ = self.env.step(action)
                trajectory.append((state, action, reward, next_state, done))
                state = next_state
            trajectories.append(trajectory)
        return trajectories

    def compute_rewards(self, trajectories):
        rewards = []
        for trajectory in trajectories:
            G = 0
            for t in range(len(trajectory) - 1, -1, -1):
                _, _, reward, _, _ = trajectory[t]
                G = reward + self.gamma * G
                trajectory[t] = (*trajectory[t], G)
            rewards.append(trajectory)
        return rewards

    def compute_advantages(self, trajectories):
        advantages = []
        for trajectory in trajectories:
            V = self.critic(np.expand_dims(np.array([state for state, _, _, _, _ in trajectory]), axis=0))
            V = V.numpy()[0]
            A = []
            for t in range(len(trajectory)):
                _, _, _, _, done = trajectory[t]
                delta = trajectory[t][2] + self.gamma * (1 - done) * V[t + 1] - V[t]
                A.append(delta)
            advantages.append(A)
        return advantages

    def update_actor(self, trajectories, advantages):
        for trajectory, advantage in zip(trajectories, advantages):
            states, actions, _, _, _ = zip(*trajectory)
            with tf.GradientTape() as tape:
                log_probs = self.actor(np.expand_dims(np.array(states), axis=0))[1]
                log_probs = log_probs.numpy()[0]
                ratios = tf.exp(log_probs - trajectory[0][1])
                loss = -tf.reduce_mean(tf.minimum(ratios * advantage, tf.clip_by_value(ratios, 1 - self.clip_range, 1 + self.clip_range) * advantage))
            grads = tape.gradient(loss, self.actor.trainable_variables)
            self.actor_optimizer.apply_gradients(zip(grads, self.actor.trainable_variables))

    def update_critic(self, trajectories):
        for trajectory in trajectories:
            states, _, rewards, _, _ = zip(*trajectory)
            with tf.GradientTape() as tape:
                V = self.critic(np.expand_dims(np.array(states), axis=0))
                loss = tf.reduce_mean((V - np.array(rewards)) ** 2)
            grads = tape.gradient(loss, self.critic.trainable_variables)
            self.critic_optimizer.apply_gradients(zip(grads, self.critic.trainable_variables))

    def train(self, iterations=100):
        for _ in range(iterations):
            trajectories = self.collect_trajectories()
            rewards = self.compute_rewards(trajectories)
            advantages = self.compute_advantages(rewards)
            self.update_actor(trajectories, advantages)
            self.update_critic(trajectories)

class Actor(tf.keras.Model):
    def __init__(self, observation_space, action_space):
        super().__init__()
        self.observation_space = observation_space
        self.action_space = action_space

        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
        self.dense2 = tf.keras.layers.Dense(64, activation='relu')
        self.actor = tf.keras.layers.Dense(self.action_space.n)

    def call(self, inputs):
        x = self.dense1(inputs)
        x = self.dense2(x)
        logits = self.actor(x)
        probs = tf.nn.softmax(logits)
        return probs, logits

class Critic(tf.keras.Model):
    def __init__(self, observation_space):
        super().__init__()
        self.observation_space = observation_space

        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
        self.dense2 = tf.keras.layers.Dense(64, activation='relu')
        self.critic = tf.keras.layers.Dense(1)

    def call(self, inputs):
        x = self.dense1(inputs)
        x = self.dense2(x)
        value = self.critic(x)
        return value