dynamics method

Map<String, Tensor> dynamics(
  1. Tensor state,
  2. int action
)

Implementation

Map<String, Tensor> dynamics(Tensor state, int action) {
  Tensor actionEmb = transformer.wte.getRow(action);
  // Move the hidden state forward through one transformer block
  Tensor nextState = transformer.blocks[0].forward(state + actionEmb, Tensor.zeros([1, embedSize]));

  // Simple reward head
  Tensor reward = nextState.matmul(transformer.vW1).matmul(transformer.vW2);
  reward.data[0] = _mathTanh(reward.data[0]);

  return {"state": nextState, "reward": reward};
}