code: def main(args: argparse.Namespace):
      run_name = f"{args.env_id}_{args.exp_name}_{args.seed}_{int(time.time())}"
      logger = TensorBoardLogger(
          root_dir=os.path.join(
              "logs",
              "fabric_logs",
              datetime.today().strftime("%Y-%m-%d_%H-%M-%S"),
          ),
          name=run_name,
      )
  
      # Initialize Fabric
      fabric = Fabric(loggers=logger)
      rank = fabric.global_rank
      world_size = fabric.world_size
      device = fabric.device
      fabric.seed_everything(args.seed)
      torch.backends.cudnn.deterministic = args.torch_deterministic
  
      # Log hyperparameters
      fabric.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 + rank * args.num_envs + i,
                  rank,
                  args.capture_video,
                  logger.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 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,
      )
      optimizer = agent.configure_optimizers(args.learning_rate)
      agent, optimizer = fabric.setup(agent, optimizer)
  
      # Player metrics
      rew_avg = torchmetrics.MeanMetric().to(device)
      ep_len_avg = torchmetrics.MeanMetric().to(device)
  
      # Local data
      obs = torch.zeros(
          (args.num_steps, args.num_envs)
          + envs.single_observation_space.shape,
          device=device,
      )
      actions = torch.zeros(
          (args.num_steps, args.num_envs)
          + envs.single_action_space.shape,
          device=device,
      )
      logprobs = torch.zeros(
          (args.num_steps, args.num_envs), device=device
      )
      rewards = torch.zeros(
          (args.num_steps, args.num_envs), device=device
      )
      dones = torch.zeros(
          (args.num_steps, args.num_envs), device=device
      )
      values = torch.zeros(
          (args.num_steps, args.num_envs), device=device
      )
  
      # Global variables
      global_step = 0
      start_time = time.time()
      single_global_rollout = int(
          args.num_envs * args.num_steps * world_size
      )
      num_updates = args.total_timesteps // single_global_rollout
  
      # 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, device=device)
      for update in range(1, num_updates + 1):
          # Learning rate annealing
          if args.anneal_lr:
              linear_annealing(
                  optimizer, update, num_updates, args.learning_rate
              )
          fabric.log(
              "Info/learning_rate",
              optimizer.param_groups[0]["lr"],
              global_step,
          )
  
          for step in range(0, args.num_steps):
              global_step += args.num_envs * world_size
              obs[step] = next_obs
              dones[step] = next_done
  
              # Sample an action given the observation received by the environment
              with torch.no_grad():
                  (
                      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 args.share_data:
              # Gather all the tensors from all the world and reshape them
              gathered_data = fabric.all_gather(local_data)
              for k, v in gathered_data.items():
                  if k == "obs":
                      gathered_data[k] = v.reshape(
                          (-1,) + envs.single_observation_space.shape
                      )
                  elif k == "actions":
                      gathered_data[k] = v.reshape(
                          (-1,) + envs.single_action_space.shape
                      )
                  else:
                      gathered_data[k] = v.reshape(-1)
          else:
              gathered_data = local_data
  
          # Train the agent
          train(
              fabric, agent, optimizer, gathered_data, global_step, args
          )
          fabric.log(
              "Time/step_per_second",
              int(global_step / (time.time() - start_time)),
              global_step,
          )
  
      envs.close()
      if fabric.is_global_zero:
          test(agent.module, device, fabric.logger.experiment, args)
  
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 2 for calling a global function with example usage: main(args) without return types