predict method

Map<String, Tensor> predict(
  1. Tensor state
)

Implementation

Map<String, Tensor> predict(Tensor state) {
  // We use the specialized MuZero forward we added to the transformer
  // If you haven't added the specific heads yet, this calls the logic:

  // 1. Policy (using your lmHead/policyHead logic)
  Tensor logits = transformer.lmHead.forward(state);

  // 2. Value (Internal MLP logic)
  Tensor vHidden = state.matmul(transformer.vW1);
  for (int i = 0; i < vHidden.data.length; i++) {
    if (vHidden.data[i] < 0) vHidden.data[i] = 0; // ReLU
  }
  Tensor value = vHidden.matmul(transformer.vW2);
  for (int i = 0; i < value.data.length; i++) {
    value.data[i] = _mathTanh(value.data[i]); // Tanh
  }

  return {"policy": logits, "value": value};
}