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 Loss
    double pLoss = 0;
    final policyLogits = pred['policy']!.data;
    double maxLogit = policyLogits.reduce(math.max);
    double sumExp = 0;
    for (var v in policyLogits)
      sumExp += math.exp((v - maxLogit).clamp(-10, 10));

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

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

    totalLoss += (pLoss + vLoss);
  }

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