docstr: Decompose rewards at the last step to individual steps using transformer modules.
  
  Args:
      nhead: the number of heads in the multiheadattention models (default=8).
      num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
      dim_feedforward: the dimension of the feedforward network model (default=2048).
      dropout: the dropout value (default=0.1).
      activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
      layer_norm_eps: the eps value in layer normalization components (default=1e-5).
docstr_tok: ["Decompose","rewards","at","the","last","step","to","individual","steps","using","transformer","modules",".","Args",":","nhead",":","the","number","of","heads","in","the","multiheadattention","models","(","default=8",")",".","num_encoder_layers",":","the","number","of","sub-encoder-layers","in","the","encoder","(","default=6",")",".","dim_feedforward",":","the","dimension","of","the","feedforward","network","model","(","default=2048",")",".","dropout",":","the","dropout","value","(","default=0.1",")",".","activation",":","the","activation","function","of","encoder\/decoder","intermediate","layer",",","relu","or","gelu","(","default=relu",")",".","layer_norm_eps",":","the","eps","value","in","layer","normalization","components","(","default=1e-5",")","."]
code: def __init__(
      self,
      state_dim: int,
      action_dim: int,
      d_model: int,
      nhead: int = 2,
      num_encoder_layers: int = 2,
      dim_feedforward: int = 128,
      dropout: float = 0.0,
      activation: str = "relu",
      last_layer_activation: str = "leaky_relu",
      layer_norm_eps: float = 1e-5,
      max_len: int = 10,
  ):
      """
      Decompose rewards at the last step to individual steps using transformer modules.
  
      Args:
          nhead: the number of heads in the multiheadattention models (default=8).
          num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
          dim_feedforward: the dimension of the feedforward network model (default=2048).
          dropout: the dropout value (default=0.1).
          activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
          layer_norm_eps: the eps value in layer normalization components (default=1e-5).
      """
      super().__init__()
  
      self.state_dim = state_dim
      self.action_dim = action_dim
      # d_model: dimension of transformer input
      self.d_model = d_model
      self.nhead = nhead
      self.num_encoder_layers = num_encoder_layers
      self.dim_feedforward = dim_feedforward
      self.dropout = dropout
      self.activation = activation
      self.layer_norm_eps = layer_norm_eps
      self.max_len = max_len
  
      # map input features to higher latent space before sending to transformer
      self.fc_in = nn.Sequential(
          nn.Linear(self.state_dim + self.action_dim, self.d_model),
          nn.ReLU(),
      )
  
      # use transformer encoder to get reward logits for each step
      encoder_layer = PETransformerEncoderLayer(
          self.d_model,
          nhead,
          dim_feedforward,
          dropout,
          activation,
          layer_norm_eps,
          max_len=self.max_len,
          batch_first=False,
      )
      self.transformer = nn.TransformerEncoder(
          encoder_layer,
          num_encoder_layers,
      )
      self.fc_out = nn.Linear(self.d_model, 1)
      self.output_activation = ACTIVATION_MAP[last_layer_activation]()
  
imports: import logging
  import math
  from typing import List
  import torch
  import torch.nn
  import torch.nn.functional
  from reagent.core import parameters, types
  from reagent.models import convolutional_network, fully_connected_network
  from reagent.models.base import ModelBase
  from reagent.models.fully_connected_network import ACTIVATION_MAP
label_desc: Use image node_id 1 to create a new TransformerSyntheticRewardNet object from inherited base classes: nn with example: obj = TransformerSyntheticRewardNet(state_dim, action_dim, d_model, nhead, num_encoder_layers, dim_feedforward, dropout, activation, last_layer_activation, layer_norm_eps, max_len)