code: def train(
      agent: PPOAgent,
      optimizer: torch.optim.Optimizer,
      data: Dict[str, Tensor],
      logger: SummaryWriter,
      global_step: int,
      args: argparse.Namespace,
  ):
      indexes = list(range(data["obs"].shape[0]))
      if args.share_data:
          sampler = DistributedSampler(
              indexes,
              num_replicas=distributed.get_world_size(),
              rank=distributed.get_rank(),
              shuffle=True,
              seed=args.seed,
          )
      else:
          sampler = RandomSampler(indexes)
      sampler = BatchSampler(
          sampler, batch_size=args.per_rank_batch_size, drop_last=False
      )
      per_epoch_losses = torch.tensor(
          [0.0, 0.0, 0.0], device=data["obs"].device
      )
      for epoch in range(args.update_epochs):
          if args.share_data:
              sampler.sampler.set_epoch(epoch)
          for batch_idxes in sampler:
              _, newlogprob, entropy, newvalue = agent(
                  data["obs"][batch_idxes],
                  data["actions"].long()[batch_idxes],
              )
              logratio = newlogprob - data["logprobs"][batch_idxes]
              ratio = logratio.exp()
  
              advantages = data["advantages"][batch_idxes]
              if args.normalize_advantages:
                  advantages = (advantages - advantages.mean()) / (
                      advantages.std() + 1e-8
                  )
  
              # Policy loss
              pg_loss = policy_loss(advantages, ratio, args.clip_coef)
              per_epoch_losses[0] += pg_loss.detach()
  
              # Value loss
              v_loss = value_loss(
                  newvalue,
                  data["values"][batch_idxes],
                  data["returns"][batch_idxes],
                  args.clip_coef,
                  args.clip_vloss,
                  args.vf_coef,
              )
              per_epoch_losses[1] += v_loss.detach()
  
              # Entropy loss
              ent_loss = entropy_loss(entropy, args.ent_coef)
              per_epoch_losses[2] += ent_loss.detach()
  
              # Overall loss
              loss = pg_loss + ent_loss + v_loss
  
              optimizer.zero_grad(set_to_none=True)
              loss.backward()
              nn.utils.clip_grad_norm_(
                  agent.parameters(), args.max_grad_norm
              )
              optimizer.step()
  
          # Log
          distributed.reduce(per_epoch_losses, dst=0)
          if logger is not None:
              per_epoch_losses = per_epoch_losses / (
                  len(sampler) * distributed.get_world_size()
              )
              logger.add_scalar(
                  "Loss/policy_loss", per_epoch_losses[0], global_step
              )
              logger.add_scalar(
                  "Loss/value_loss", per_epoch_losses[1], global_step
              )
              logger.add_scalar(
                  "Loss/entropy_loss", per_epoch_losses[2], global_step
              )
          per_epoch_losses.fill_(0)
  
file_path: pytorch-lightning/examples/fabric/reinforcement_learning/train_torch.py
imports: import argparse
  import os
  import random
  import time
  from datetime import datetime
  from typing import Dict
  import gymnasium
  import torch
  import torch.distributed
  import torch.nn
  import torch.optim
  from rl.agent import PPOAgent
  from rl.loss import entropy_loss, policy_loss, value_loss
  from rl.utils import linear_annealing, make_env, parse_args, test
  from torch import Tensor
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import BatchSampler, DistributedSampler, RandomSampler
  from torch.utils.tensorboard import SummaryWriter
label_desc: Use image node_id 1 for calling a global function with example usage: train(agent, optimizer, data, logger, global_step, args) without return types