code: def player(
      args,
      world_collective: TorchCollective,
      player_trainer_collective: TorchCollective,
  ):
      run_name = f"{args.env_id}_{args.exp_name}_{args.seed}"
      logger = TensorBoardLogger(
          root_dir=os.path.join(
              "logs",
              "fabric_decoupled_logs",
              datetime.today().strftime("%Y-%m-%d_%H-%M-%S"),
          ),
          name=run_name,
      )
      log_dir = logger.log_dir
  
      # Initialize Fabric object
      fabric = Fabric(
          loggers=logger,
          accelerator="cuda" if args.player_on_gpu else "cpu",
      )
      device = fabric.device
      fabric.seed_everything(args.seed)
      torch.backends.cudnn.deterministic = args.torch_deterministic
  
      # Log hyperparameters
      logger.experiment.add_text(
          "hyperparameters",
          "|param|value|\n|-|-|\n%s"
          % (
              "\n".join(
                  [
                      f"|{key}|{value}|"
                      for key, value in vars(args).items()
                  ]
              )
          ),
      )
  
      # Environment setup
      envs = gym.vector.SyncVectorEnv(
          [
              make_env(
                  args.env_id,
                  args.seed + i,
                  0,
                  args.capture_video,
                  log_dir,
                  "train",
              )
              for i in range(args.num_envs)
          ]
      )
      assert isinstance(
          envs.single_action_space, gym.spaces.Discrete
      ), "only discrete action space is supported"
  
      # Define the agent
      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,
      ).to(device)
      flattened_parameters = torch.empty_like(
          torch.nn.utils.convert_parameters.parameters_to_vector(
              agent.parameters()
          ),
          device=device,
      )
  
      # Receive the first weights from the rank-1, a.k.a. the first of the trainers
      # In this way we are sure that before the first iteration everyone starts with the same parameters
      player_trainer_collective.broadcast(flattened_parameters, src=1)
      torch.nn.utils.convert_parameters.vector_to_parameters(
          flattened_parameters, agent.parameters()
      )
  
      # Player metrics
      rew_avg = MeanMetric(sync_on_compute=False).to(device)
      ep_len_avg = MeanMetric(sync_on_compute=False).to(device)
  
      # Local data
      obs = torch.zeros(
          (args.num_steps, args.num_envs)
          + envs.single_observation_space.shape
      ).to(device)
      actions = torch.zeros(
          (args.num_steps, args.num_envs)
          + envs.single_action_space.shape
      ).to(device)
      logprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)
      rewards = torch.zeros((args.num_steps, args.num_envs)).to(device)
      dones = torch.zeros((args.num_steps, args.num_envs)).to(device)
      values = torch.zeros((args.num_steps, args.num_envs)).to(device)
  
      # Global variables
      global_step = 0
      start_time = time.time()
      single_global_step = int(args.num_envs * args.num_steps)
      num_updates = args.total_timesteps // single_global_step
      if not args.share_data:
          if single_global_step < world_collective.world_size - 1:
              raise RuntimeError(
                  "The number of trainers ({}) is greater than the available collected data ({}). ".format(
                      world_collective.world_size - 1,
                      single_global_step,
                  )
                  + "Consider to lower the number of trainers at least to the size of available collected data"
              )
          chunks_sizes = [
              len(chunk)
              for chunk in torch.tensor_split(
                  torch.arange(single_global_step),
                  world_collective.world_size - 1,
              )
          ]
  
      # Broadcast num_updates to all the world
      update_t = torch.tensor(
          [num_updates], device=device, dtype=torch.float32
      )
      world_collective.broadcast(update_t, src=0)
  
      # Get the first environment observation and start the optimization
      next_obs = torch.tensor(
          envs.reset(seed=args.seed)[0], device=device
      )
      next_done = torch.zeros(args.num_envs).to(device)
      for _ in range(1, num_updates + 1):
          for step in range(0, args.num_steps):
              global_step += args.num_envs
              obs[step] = next_obs
              dones[step] = next_done
  
              # Sample an action given the observation received by the environment
              action, logprob, _, value = agent.get_action_and_value(
                  next_obs
              )
              values[step] = value.flatten()
              actions[step] = action
              logprobs[step] = logprob
  
              # Single environment step
              next_obs, reward, done, truncated, info = envs.step(
                  action.cpu().numpy()
              )
              done = torch.logical_or(
                  torch.tensor(done), torch.tensor(truncated)
              )
              rewards[step] = torch.tensor(reward, device=device).view(
                  -1
              )
              next_obs, next_done = torch.tensor(
                  next_obs, device=device
              ), done.to(device)
  
              if "final_info" in info:
                  for i, agent_final_info in enumerate(
                      info["final_info"]
                  ):
                      if (
                          agent_final_info is not None
                          and "episode" in agent_final_info
                      ):
                          fabric.print(
                              f"Rank-0: global_step={global_step}, reward_env_{i}={agent_final_info['episode']['r'][0]}"
                          )
                          rew_avg(agent_final_info["episode"]["r"][0])
                          ep_len_avg(
                              agent_final_info["episode"]["l"][0]
                          )
  
          # Sync the metrics
          rew_avg_reduced = rew_avg.compute()
          if not rew_avg_reduced.isnan():
              fabric.log(
                  "Rewards/rew_avg", rew_avg_reduced, global_step
              )
          ep_len_avg_reduced = ep_len_avg.compute()
          if not ep_len_avg_reduced.isnan():
              fabric.log(
                  "Game/ep_len_avg", ep_len_avg_reduced, global_step
              )
          rew_avg.reset()
          ep_len_avg.reset()
  
          # Estimate returns with GAE (https://arxiv.org/abs/1506.02438)
          returns, advantages = agent.estimate_returns_and_advantages(
              rewards,
              values,
              dones,
              next_obs,
              next_done,
              args.num_steps,
              args.gamma,
              args.gae_lambda,
          )
  
          # Flatten the batch
          local_data = {
              "obs": obs.reshape(
                  (-1,) + envs.single_observation_space.shape
              ),
              "logprobs": logprobs.reshape(-1),
              "actions": actions.reshape(
                  (-1,) + envs.single_action_space.shape
              ),
              "advantages": advantages.reshape(-1),
              "returns": returns.reshape(-1),
              "values": values.reshape(-1),
          }
          if not args.player_on_gpu and args.cuda:
              for v in local_data.values():
                  v = v.pin_memory()
  
          # Send data to the training agents
          if args.share_data:
              world_collective.broadcast_object_list(
                  [local_data], src=0
              )
          else:
              # Split data in an even way, when possible
              perm = torch.randperm(single_global_step, device=device)
              chunks = [
                  {} for _ in range(world_collective.world_size - 1)
              ]
              for k, v in local_data.items():
                  chunked_local_data = v[perm].split(chunks_sizes)
                  for i in range(len(chunks)):
                      chunks[i][k] = chunked_local_data[i]
  
              world_collective.scatter_object_list(
                  [None], [None] + chunks, src=0
              )
  
          # Gather metrics from the trainers to be plotted
          metrics = [None]
          player_trainer_collective.broadcast_object_list(
              metrics, src=1
          )
  
          # Wait the trainers to finish
          player_trainer_collective.broadcast(
              flattened_parameters, src=1
          )
  
          # Convert back the parameters
          torch.nn.utils.convert_parameters.vector_to_parameters(
              flattened_parameters, agent.parameters()
          )
  
          fabric.log_dict(metrics[0], global_step)
          fabric.log_dict(
              {
                  "Time/step_per_second": int(
                      global_step / (time.time() - start_time)
                  )
              },
              global_step,
          )
  
      if args.share_data:
          world_collective.broadcast_object_list([-1], src=0)
      else:
          world_collective.scatter_object_list(
              [None],
              [None] + [-1] * (world_collective.world_size - 1),
              src=0,
          )
      envs.close()
      test(agent, device, fabric.logger.experiment, args)
  
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 1 for calling a global function with example usage: player(args, world_collective, player_trainer_collective) without return types