trainStep method

double trainStep({ int batchSize = 32, })

Implementation

double trainStep({int batchSize = 32}) {
  if (replayBuffer.length < batchSize) return 0.0;
  optimizer.zeroGrad();
  double totalLoss = 0;
  final random = math.Random();

  for (int i = 0; i < batchSize; i++) {
    final sample = replayBuffer[random.nextInt(replayBuffer.length)];
    final state = model.represent(sample.observations);
    final pred = model.predict(state);

    // Policy Cross-Entropy
    double pLoss = 0;
    final policy = pred['policy']!.data;
    double maxLogit = policy.reduce(math.max);
    double sumExp = 0;
    for (var v in policy) sumExp += math.exp((v - maxLogit).clamp(-10, 10));

    for (int j = 0; j < 4098; j++) {
      if (sample.targetPi[j] > 0) {
        double logSoftmax = (policy[j] - maxLogit) - math.log(sumExp + 1e-10);
        pLoss -= sample.targetPi[j] * logSoftmax;
      }
    }

    // Value MSE
    double vErr = pred['value']!.data[0] - sample.targetValue!;
    double vLoss = vErr * vErr;

    totalLoss += (pLoss + 0.5 * vLoss);
  }

  optimizer.step();
  return totalLoss / batchSize;
}