javascript - 贝尔曼方程在 TFjs 中的损失

Question

我正在尝试使用 p5 和 tfjs 在 JS 中实现一个简单的 Pong Q-Network。为了训练网络，我首先创建了一个自定义损失函数，其中我传递了一个标签张量，该张量仅包含动作 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). 请确保使用变量的操作在传递给最小化（）的函数 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();
}

标签： javascriptartificial-intelligencereinforcement-learningtensorflow.js