code: def create_data(
      state_dim,
      action_dim,
      seq_len,
      batch_size,
      num_batches,
      binary_reward=False,
  ):
      SCALE = 2
      # reward is a linear function of (state, action)
      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))
  
          # reward_matrix shape: batch_size x seq_len
          reward_matrix = torch.matmul(
              torch.cat((state, action), dim=2), weight
          ).transpose(0, 1)
          if binary_reward:
              reward_matrix = torch.sigmoid(reward_matrix)
          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 1 for calling a global function with example usage: create_data(state_dim, action_dim, seq_len, batch_size, num_batches, binary_reward) and returns: weight,  data