Как преобразовать SymbolicTensor в np.ndarray при использовании @tf.function в Keras для RL агента?

Я хочу написать RL агента на keras, но у меня возникает проблема. Класс я написал, но при попытке вызвать функцию DQN.train_step вылетает ошибка AttributeError: 'SymbolicTensor' object has no attribute 'numpy'. Из-за декоратора @tf.function нейросеть возвращает объект класса tensorflow.python.framework.ops.SymbolicTensor, а не как обычно np.ndarray. Я не хочу убирать декоратор, потому что он ускоряет работу программы примерно в 50 раз. Проблема в этом участке кода:

with tf.GradientTape() as tape:
    current_q_values = self.predict(state)
    next_q_values = self.predict(next_state)
    max_next_q = tf.reduce_max(next_q_values)
    target_q_values = current_q_values.numpy()  # здесь возникает ошибка
    target_q_values[0, action] = reward + self.discount_factor * max_next_q * (1 - done)
    loss = self.loss_fn(current_q_values, target_q_values)

В строке target_q_values = current_q_values.numpy() надо заменить current_q_values.numpy() на какую-то функцию или что-то ещё, что будет изменять тип данных c tensorflow.python.framework.ops.SymbolicTensor на np.ndarray, но как это сделать я не знаю.

import numpy as np
import tensorflow as tf
from keras.models import Model

class DQN:
    def __init__(self, model, optimizer, loss_fn,
                 discount_factor: int | float = 0.98, lambda_: int | float = 0):
        """
        :param model: Нейросеть, keras.models.Sequential.
        :param optimizer: Оптимизатор нейросети.
        :param loss_fn: Функция потерь нейросети.
        :param discount_factor: Дисконтирующий фактор / насколько важна награда в следующем состоянии для агента.
        :param lambda_: Насколько важно помнить ранее приобретённые навыки.
        """

        self.model = layers
        self.optimizer = optimizer
        self.loss_fn = loss_fn
        self.discount_factor = discount_factor
        self.lambda_ = lambda_
        self.prior_weights = self.fisher_matrix = None

    @tf.function
    def train_step(self, state: np.ndarray, action: int, reward: int | float,
                       next_state: np.ndarray, done: int):
        """
        :param state: Состояние, на основе которого агент принял решение совершить действие action.
        :param action: Действие, которое совершило переход из состояния state в состояние next_state.
        :param reward: Награда за действие action.
        :param next_state: Состояние, в которое был совершён переход из состояния state с помощью действия action.
        :param done: Терминальное состояние, принимает значение 0 или 1.
        """

        with tf.GradientTape() as tape:
            current_q_values = self.predict(state)
            next_q_values = self.predict(next_state)
            max_next_q = tf.reduce_max(next_q_values)
            target_q_values = current_q_values.numpy()
            target_q_values[0, action] = reward + self.discount_factor * max_next_q * (1 - done)
            loss = self.loss_fn(current_q_values, target_q_values)

            loss += compute_penalty_loss(self.weights, self.prior_weights, self.fisher_matrix, self.lambda_)

        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

    def predict(self, x, *args, **kwargs):
        return self.model(x, *args, **kwargs)

Я не знаю как мне решить эту проблему, помогите мне, пожалуйста.

Версии библиотек, которые я использую:

• tensorflow 2.16.2
• keras 3.4.1
• numpy 1.26.4