code: def _action_docs(
      self,
      state: rlt.FeatureData,
      action: torch.Tensor,
      terminal_mask: Optional[torch.Tensor] = None,
  ) -> rlt.DocList:
      # for invalid indices, simply set action to 0 so we can batch index still
      if terminal_mask is not None:
          assert terminal_mask.shape == (
              action.shape[0],
          ), f"{terminal_mask.shape} != 0th dim of {action.shape}"
          action[terminal_mask] = torch.zeros_like(
              action[terminal_mask]
          )
      docs = state.candidate_docs
      assert docs is not None
      return docs.select_slate(action)
  
imports: import enum
  import logging
  from typing import Optional
  import reagent.core.parameters
  import reagent.core.types
  import torch
  import torch.nn.functional
  from reagent.core.dataclasses import field
  from reagent.optimizer import Optimizer__Union, SoftUpdate
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.rl_trainer_pytorch import RLTrainerMixin
label_desc: Use image node_id 3 for calling the SlateQTrainer obj's underlying member method code with example usage: obj._action_docs(state, action, terminal_mask) and returns: docs