code: def get_Q(
      seq2reward_network: Seq2RewardNetwork,
      cur_state: torch.Tensor,
      all_permut: torch.Tensor,
  ) -> torch.Tensor:
      """
      Input:
          cur_state: the current state from where we start planning.
              shape: batch_size x state_dim
          all_permut: all action sequences (sorted in lexical order) for enumeration
              shape: seq_len x num_perm x action_dim
      """
      batch_size = cur_state.shape[0]
      _, num_permut, num_action = all_permut.shape
      num_permut_per_action = int(num_permut / num_action)
  
      preprocessed_state = cur_state.unsqueeze(0).repeat_interleave(
          num_permut, dim=1
      )
      state_feature_vector = rlt.FeatureData(preprocessed_state)
  
      # expand action to match the expanded state sequence
      action = rlt.FeatureData(all_permut.repeat(1, batch_size, 1))
      acc_reward = seq2reward_network(
          state_feature_vector, action
      ).acc_reward.reshape(
          batch_size, num_action, num_permut_per_action
      )
  
      # The permuations are generated with lexical order
      # the output has shape [num_perm, num_action,1]
      # that means we can aggregate on the max reward
      # then reshape it to (BATCH_SIZE, ACT_DIM)
      max_acc_reward = (
          torch.max(acc_reward, dim=2)
          .values.detach()
          .reshape(batch_size, num_action)
      )
  
      return max_acc_reward
  
imports: import logging
  import reagent.core.types
  import torch
  import torch.nn
  import torch.nn.functional
  from reagent.core.parameters import Seq2RewardTrainerParameters
  from reagent.models.fully_connected_network import FullyConnectedNetwork
  from reagent.models.seq2reward_model import Seq2RewardNetwork
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.utils import gen_permutations
label_desc: Use image node_id 2 for calling a global function with example usage: get_Q(seq2reward_network, cur_state, all_permut) and returns: max_acc_reward