TFjs中Bellman方程的损失

Question

我正在尝试使用p5和tfjs在JS中实现一个简单的Pong Q网络。 为了训练网络，我首先创建了一个自定义损失函数，我在其中传递了一个标签张量，该张量仅包含动作a_t

的标签

TFjs似乎并不真正喜欢形状与模型预测值不同的事实，因此我设计了另一个自定义损失函数，其中标签输入为形状[batchSize, 3]的张量（最多3个动作） （向下且无），其中每个元素的格式为[0,0,y_j]或[0,y_j,0]或[y_j,0,0]（y_j在我应该与预测张量进行比较的位置，其他地方为0）。在这里：

function bellmanLoss(predictions, labels)
{
    let predictions_buffer = predictions.buffer();
    let labels_buffer = labels.buffer();
    let length = labels.shape[0];
    predictions.dispose();
    labels.dispose();
    let loss = 0;
    for(let i = 0; i < length; i++)
    {
        for(let j = 0; j < 3; j++)
        {
            if(labels_buffer.get(i,j) != 0)
            {
                loss += Math.pow(labels_buffer.get(i,j) - predictions_buffer.get(i,j), 2);
                break;
            }
        }
    }
    return tf.tensor(loss);
}

但是这里我是从零开始构造张量的，所以我使用model.fit: "Error: Cannot find a connection between any variable and the result of the loss function y=f(x)得到了这个错误。请确保使用变量的操作在传递给minimum（）的函数f中。 有没有办法使损失与model.fit兼容，还是我必须手动调整模型权重（这很痛苦）？

更新：我进行了一些更改，使其更具“张力”，而且似乎朝着正确的方向发展：

function bellmanLoss(preds, labels)
{
    let mask = tf.cast(labels, 'bool');
    let zeros = tf.zerosLike(preds);
    let clean_preds = preds.where(mask, zeros);
    return tf.squaredDifference(clean_preds, labels).mean();
}

但是，我仍然需要找到另一种选择，因为

  

“错误：无法计算梯度：在何处找不到梯度函数”

最终更新：我找到了一种不使用tf.where的方法，并且可以正常工作

function bellmanLoss(preds, labels)
{
    let mask_b = tf.cast(labels, 'bool');
    let mask = tf.cast(mask_b, 'float32');
    let clean_preds = preds.mul(mask);
    return tf.squaredDifference(clean_preds, labels).mean();
}