code: def create_sequence_data(
      state_dim, action_dim, seq_len, batch_size, num_batches
  ):
      SCALE = 2
      weight = SCALE * torch.randn(state_dim + action_dim)
  
      data = [None for _ in range(num_batches)]
  
      for i in range(num_batches):
          state = SCALE * torch.randn(seq_len, batch_size, state_dim)
          action = SCALE * torch.randn(seq_len, batch_size, action_dim)
          # random valid step
          valid_step = torch.randint(1, seq_len + 1, (batch_size, 1))
  
          feature_mask = torch.arange(seq_len).repeat(batch_size, 1)
          feature_mask = (
              feature_mask >= (seq_len - valid_step)
          ).float()
          assert feature_mask.shape == (
              batch_size,
              seq_len,
          ), feature_mask.shape
          feature_mask = feature_mask.transpose(0, 1).unsqueeze(-1)
          assert feature_mask.shape == (
              seq_len,
              batch_size,
              1,
          ), feature_mask.shape
  
          feature = torch.cat((state, action), dim=2)
          masked_feature = feature * feature_mask
  
          # seq_len, batch_size, state_dim + action_dim
          left_shifted = torch.cat(
              (
                  masked_feature.narrow(0, 1, seq_len - 1),
                  torch.zeros(1, batch_size, state_dim + action_dim),
              ),
              dim=0,
          )
          # seq_len, batch_size, state_dim + action_dim
          right_shifted = torch.cat(
              (
                  torch.zeros(1, batch_size, state_dim + action_dim),
                  masked_feature.narrow(0, 0, seq_len - 1),
              ),
              dim=0,
          )
          # reward_matrix shape: batch_size x seq_len
          reward_matrix = torch.matmul(
              left_shifted + right_shifted, weight
          ).transpose(0, 1)
  
          mask = torch.arange(seq_len).repeat(batch_size, 1)
          mask = (mask >= (seq_len - valid_step)).float()
          reward = (reward_matrix * mask).sum(dim=1).reshape(-1, 1)
  
          data[i] = rlt.MemoryNetworkInput(
              state=rlt.FeatureData(state),
              action=rlt.FeatureData(action),
              valid_step=valid_step,
              reward=reward,
              # the rest fields will not be used
              next_state=torch.tensor([]),
              step=torch.tensor([]),
              not_terminal=torch.tensor([]),
              time_diff=torch.tensor([]),
          )
  
      return weight, data
  
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 2 for calling a global function with example usage: create_sequence_data(state_dim, action_dim, seq_len, batch_size, num_batches) and returns: weight,  data