code: def train(
      fabric: Fabric,
      agent: PPOLightningAgent,
      optimizer: torch.optim.Optimizer,
      data: Dict[str, Tensor],
      global_step: int,
      args: argparse.Namespace,
  ):
      indexes = list(range(data["obs"].shape[0]))
      if args.share_data:
          sampler = DistributedSampler(
              indexes,
              num_replicas=fabric.world_size,
              rank=fabric.global_rank,
              shuffle=True,
              seed=args.seed,
          )
      else:
          sampler = RandomSampler(indexes)
      sampler = BatchSampler(
          sampler, batch_size=args.per_rank_batch_size, drop_last=False
      )
  
      for epoch in range(args.update_epochs):
          if args.share_data:
              sampler.sampler.set_epoch(epoch)
          for batch_idxes in sampler:
              loss = agent.training_step(
                  {k: v[batch_idxes] for k, v in data.items()}
              )
              optimizer.zero_grad(set_to_none=True)
              fabric.backward(loss)
              fabric.clip_gradients(
                  agent, optimizer, max_norm=args.max_grad_norm
              )
              optimizer.step()
          agent.on_train_epoch_end(global_step)
  
file_path: pytorch-lightning/examples/fabric/reinforcement_learning/train_fabric.py
imports: import argparse
  import os
  import time
  from datetime import datetime
  from typing import Dict
  import gymnasium
  import torch
  import torchmetrics
  from lightning.fabric import Fabric
  from lightning.fabric.loggers import TensorBoardLogger
  from rl.agent import PPOLightningAgent
  from rl.utils import linear_annealing, make_env, parse_args, test
  from torch import Tensor
  from torch.utils.data import BatchSampler, DistributedSampler, RandomSampler
label_desc: Use image node_id 1 for calling a global function with example usage: train(fabric, agent, optimizer, data, global_step, args) without return types