dynamics method

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

Implementation

Map<String, Tensor> dynamics(Tensor state, int action) {
  // FIX: Changed tokenEmbedding to wte
  Tensor actionEmb = transformer.wte.getRow(action);

  // Transition: Add the action embedding to the current state
  // and pass through one Transformer Block to 'process' the move
  Tensor nextState = transformer.blocks[0].forward(
    state + actionEmb,
    Tensor.zeros([1, embedSize]),
  );

  return {
    "state": nextState,
    "reward": Tensor.fill([1], 0.0),
  };
}