docstr: Reward at each step is a linear function of states and actions in a
  context window around the step.
  
  However, we can only observe aggregated reward at the last step
docstr_tok: ["Reward","at","each","step","is","a","linear","function","of","states","and","actions","in","a","context","window","around","the","step",".","However",",","we","can","only","observe","aggregated","reward","at","the","last","step"]
code: def test_transformer_parametric_reward(self):
      """
      Reward at each step is a linear function of states and actions in a
      context window around the step.
  
      However, we can only observe aggregated reward at the last step
      """
      state_dim = 10
      action_dim = 2
      seq_len = 5
      batch_size = 512
      num_batches = 10000
      d_model = 64
      nhead = 8
      num_encoder_layers = 1
      dim_feedforward = 64
      last_layer_activation = "linear"
      max_len = seq_len + 1
      reward_net = SyntheticRewardNet(
          TransformerSyntheticRewardNet(
              state_dim=state_dim,
              action_dim=action_dim,
              d_model=d_model,
              nhead=nhead,
              num_encoder_layers=num_encoder_layers,
              dim_feedforward=dim_feedforward,
              dropout=0.0,
              activation="relu",
              last_layer_activation=last_layer_activation,
              layer_norm_eps=1e-5,
              max_len=max_len,
          )
      )
      optimizer = Optimizer__Union(Adam=classes["Adam"]())
      trainer = RewardNetTrainer(reward_net, optimizer)
      trainer.set_reporter(
          RewardNetworkReporter(
              trainer.loss_type,
              str(reward_net),
          )
      )
      weight, data = create_sequence_data(
          state_dim, action_dim, seq_len, batch_size, num_batches
      )
  
      threshold = 0.25
      avg_eval_loss = train_and_eval(trainer, data)
      assert (
          avg_eval_loss < threshold
      ), "loss = {:.4f} larger than threshold {}".format(
          avg_eval_loss, threshold
      )
  
imports: import logging
  import unittest
  import pytorch_lightning
  import torch
  from reagent.core import parameters, types
  from reagent.models.synthetic_reward import NGramConvolutionalNetwork, NGramFullyConnectedNetwork, SequenceSyntheticRewardNet, SingleStepSyntheticRewardNet, SyntheticRewardNet, TransformerSyntheticRewardNet
  from reagent.optimizer.union import classes, Optimizer__Union
  from reagent.reporting.reward_network_reporter import RewardNetworkReporter
  from reagent.training import RewardNetTrainer
  from reagent.training.reward_network_trainer import LossFunction
  from torch.utils.data import DataLoader
label_desc: Use image node_id 9 for calling the TestSyntheticRewardTraining obj's underlying member method code with example usage: obj.test_transformer_parametric_reward() without return types