code: def trainer(
      args,
      world_collective: TorchCollective,
      player_trainer_collective: TorchCollective,
      optimization_pg: CollectibleGroup,
  ):
      global_rank = world_collective.rank
      group_rank = global_rank - 1
      group_world_size = world_collective.world_size - 1
  
      # Initialize Fabric
      fabric = Fabric(
          strategy=DDPStrategy(process_group=optimization_pg),
          accelerator="cuda" if args.cuda else "cpu",
      )
      device = fabric.device
      fabric.seed_everything(args.seed)
      torch.backends.cudnn.deterministic = args.torch_deterministic
  
      # Environment setup
      envs = gym.vector.SyncVectorEnv(
          [make_env(args.env_id, 0, 0, False, None)]
      )
      assert isinstance(
          envs.single_action_space, gym.spaces.Discrete
      ), "only discrete action space is supported"
  
      # Define the agent and the optimizer and setup them with Fabric
      agent: PPOLightningAgent = PPOLightningAgent(
          envs,
          act_fun=args.activation_function,
          vf_coef=args.vf_coef,
          ent_coef=args.ent_coef,
          clip_coef=args.clip_coef,
          clip_vloss=args.clip_vloss,
          ortho_init=args.ortho_init,
          normalize_advantages=args.normalize_advantages,
          process_group=optimization_pg,
      )
      optimizer = agent.configure_optimizers(args.learning_rate)
      agent, optimizer = fabric.setup(agent, optimizer)
  
      # Send weights to rank-0, a.k.a. the player
      if global_rank == 1:
          player_trainer_collective.broadcast(
              torch.nn.utils.convert_parameters.parameters_to_vector(
                  agent.parameters()
              ),
              src=1,
          )
  
      # Receive maximum number of updates from the player
      update = 0
      num_updates = torch.zeros(1, device=device)
      world_collective.broadcast(num_updates, src=0)
      num_updates = num_updates.item()
  
      # Start training
      while True:
          # Wait for data
          data = [None]
          if args.share_data:
              world_collective.broadcast_object_list(data, src=0)
          else:
              world_collective.scatter_object_list(
                  data,
                  [None for _ in range(world_collective.world_size)],
                  src=0,
              )
          data = data[0]
          if data == -1:
              return
  
          # Metrics dict to be sent to the player
          if group_rank == 0:
              metrics = {}
  
          # Lerning rate annealing
          if args.anneal_lr:
              linear_annealing(
                  optimizer, update, num_updates, args.learning_rate
              )
          if group_rank == 0:
              metrics["Info/learning_rate"] = optimizer.param_groups[0][
                  "lr"
              ]
          update += 1
  
          indexes = list(range(data["obs"].shape[0]))
          if args.share_data:
              sampler = DistributedSampler(
                  indexes,
                  num_replicas=group_world_size,
                  rank=group_rank,
                  shuffle=True,
                  seed=args.seed,
                  drop_last=False,
              )
          else:
              sampler = RandomSampler(indexes)
          sampler = BatchSampler(
              sampler,
              batch_size=args.per_rank_batch_size,
              drop_last=False,
          )
  
          # The Join context is needed because there can be the possibility
          # that some ranks receive less data
          with Join(
              [agent._forward_module]
          ) if not args.share_data else nullcontext():
              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].to(device)
                              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()
  
          # Sync metrics
          avg_pg_loss = agent.avg_pg_loss.compute()
          avg_value_loss = agent.avg_value_loss.compute()
          avg_ent_loss = agent.avg_ent_loss.compute()
          agent.reset_metrics()
  
          # Send updated weights to the player
          if global_rank == 1:
              metrics["Loss/policy_loss"] = avg_pg_loss
              metrics["Loss/value_loss"] = avg_value_loss
              metrics["Loss/entropy_loss"] = avg_ent_loss
              player_trainer_collective.broadcast_object_list(
                  [metrics], src=1
              )  # Broadcast metrics: fake send with object list between rank-0 and rank-1
              player_trainer_collective.broadcast(
                  torch.nn.utils.convert_parameters.parameters_to_vector(
                      agent.parameters()
                  ),
                  src=1,
              )
  
file_path: pytorch-lightning/examples/fabric/reinforcement_learning/train_fabric_decoupled.py
imports: import argparse
  import os
  import time
  from contextlib import nullcontext
  from datetime import datetime
  import gymnasium
  import torch
  from lightning.fabric import Fabric
  from lightning.fabric.loggers import TensorBoardLogger
  from lightning.fabric.plugins.collectives import TorchCollective
  from lightning.fabric.plugins.collectives.collective import CollectibleGroup
  from lightning.fabric.strategies import DDPStrategy
  from rl.agent import PPOLightningAgent
  from rl.utils import linear_annealing, make_env, parse_args, test
  from torch.distributed.algorithms.join import Join
  from torch.utils.data import BatchSampler, DistributedSampler, RandomSampler
  from torchmetrics import MeanMetric
label_desc: Use image node_id 2 for calling a global function with example usage: trainer(args, world_collective, player_trainer_collective, optimization_pg) without return types