code: def _gen_mask(
      valid_step: torch.Tensor, batch_size: int, seq_len: int
  ):
      """
      Mask for dealing with different lengths of MDPs
  
      Example:
      valid_step = [[1], [2], [3]], batch_size=3, seq_len = 4
      mask = [
          [0, 0, 0, 1],
          [0, 0, 1, 1],
          [0, 1, 1, 1],
      ]
      """
      assert valid_step.shape == (batch_size, 1)
      assert (1 <= valid_step).all()
      assert (valid_step <= seq_len).all()
      device = valid_step.device
      mask = torch.arange(seq_len, device=device).repeat(batch_size, 1)
      mask = (mask >= (seq_len - valid_step)).float()
      return mask
  
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 3 for calling a global function with example usage: _gen_mask(valid_step, batch_size, seq_len) and returns: mask