code: def one_hot_actions(
      num_actions: int,
      action: torch.Tensor,
      next_action: torch.Tensor,
      terminal: torch.Tensor,
  ):
      """
      One-hot encode actions and non-terminal next actions.
      Input shape is (batch_size, 1). Output shape is (batch_size, num_actions)
      """
      assert (
          len(action.shape) == 2
          and action.shape[1] == 1
          and next_action.shape == action.shape
      ), (
          f"Must be action with stack_size = 1, but "
          f"got shapes {action.shape}, {next_action.shape}"
      )
      action = F.one_hot(action, num_actions).squeeze(1).float()
      # next action is garbage for terminal transitions (so just zero them)
      next_action_res = torch.zeros_like(action)
      non_terminal_indices = (terminal == 0).squeeze(1)
      next_action_res[non_terminal_indices] = (
          F.one_hot(next_action[non_terminal_indices], num_actions)
          .squeeze(1)
          .float()
      )
      return action, next_action_res
  
imports: import inspect
  import logging
  from typing import Dict, Optional
  import gym
  import numpy
  import reagent.core.types
  import torch
  import torch.nn.functional
  from reagent.core.parameters import CONTINUOUS_TRAINING_ACTION_RANGE
  from reagent.gym.types import Trajectory
  from reagent.preprocessing.types import InputColumn
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.utils import rescale_actions
label_desc: Use image node_id 4 for calling a global function with example usage: one_hot_actions(num_actions, action, next_action, terminal) and returns: action,  next_action_res