code: def main(args: argparse.Namespace):
      # Init distributed environment
      is_cuda_available = torch.cuda.is_available() and args.cuda
      backend = "nccl" if is_cuda_available else "gloo"
      local_rank = int(os.environ["LOCAL_RANK"])
      global_rank = int(os.environ["RANK"])
      world_size = int(os.environ["WORLD_SIZE"])
      if is_cuda_available:
          torch.cuda.set_device(local_rank)
      device = torch.device(
          f"cuda:{local_rank}" if is_cuda_available else "cpu"
      )
      distributed.init_process_group(backend=backend)
  
      # Seed everything
      random.seed(args.seed)
      torch.manual_seed(args.seed)
      torch.cuda.manual_seed_all(args.seed)
      torch.backends.cudnn.deterministic = args.torch_deterministic
  
      # Logger
      log_dir = None
      logger = None
      run_name = f"{args.env_id}_{args.exp_name}_{args.seed}_{int(time.time())}"
      if global_rank == 0:
          log_dir = os.path.join(
              "logs",
              "torch_logs",
              datetime.today().strftime("%Y-%m-%d_%H-%M-%S"),
              run_name,
          )
          logger = SummaryWriter(log_dir=log_dir)
  
      # Log hyperparameters
      if global_rank == 0:
          logger.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 + global_rank * args.num_envs + i,
                  global_rank,
                  args.capture_video,
                  logger.log_dir if global_rank == 0 else None,
                  "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 DistributedDataParallel
      agent: PPOAgent = PPOAgent(
          envs,
          act_fun=args.activation_function,
          ortho_init=args.ortho_init,
      ).to(device)
      agent = DistributedDataParallel(
          agent,
          device_ids=[local_rank] if is_cuda_available else None,
          output_device=local_rank if is_cuda_available else None,
      )
      optimizer = optim.Adam(
          agent.parameters(), lr=args.learning_rate, eps=1e-4
      )
  
      # 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
      )
      local_rew = 0.0
      local_ep_len = 0.0
      local_num_episodes = 0.0
  
      # Global variables
      global_step = 0
      start_time = time.time()
      single_global_step = int(
          args.num_envs * args.num_steps * world_size
      )
      num_updates = args.total_timesteps // single_global_step
  
      # 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
              )
          if global_rank == 0:
              logger.add_scalar(
                  "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(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
                      ):
                          if global_rank == 0:
                              print(
                                  f"Rank-0: global_step={global_step}, "
                                  f"reward_env_{i}={agent_final_info['episode']['r'][0]}"
                              )
                          local_num_episodes += 1
                          local_rew += agent_final_info["episode"]["r"][
                              0
                          ]
                          local_ep_len += agent_final_info["episode"][
                              "l"
                          ][0]
  
          # Sync the metrics
          global_stats = torch.tensor(
              [local_rew, local_ep_len, local_num_episodes],
              device=device,
              dtype=torch.float32,
          )
          distributed.reduce(global_stats, dst=0)
          if global_rank == 0 and global_stats[2] != 0.0:
              logger.add_scalar(
                  "Rewards/rew_avg",
                  global_stats[0] / global_stats[2],
                  global_step,
              )
              logger.add_scalar(
                  "Game/ep_len_avg",
                  global_stats[1] / global_stats[2],
                  global_step,
              )
  
          # Reset metrics
          local_rew = 0
          local_ep_len = 0
          local_num_episodes = 0
  
          # Estimate returns with GAE (https://arxiv.org/abs/1506.02438)
          (
              returns,
              advantages,
          ) = agent.module.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, concat and reshape them
              gathered_data = [None for _ in range(world_size)]
              distributed.all_gather_object(gathered_data, local_data)
              processed_gathered_data = gathered_data[0]
              for i in range(1, len(gathered_data)):
                  for k in processed_gathered_data:
                      processed_gathered_data[k] = torch.cat(
                          (
                              processed_gathered_data[k].to(device),
                              gathered_data[i][k].to(device),
                          ),
                          dim=0,
                      )
          else:
              processed_gathered_data = local_data
  
          # Train the agent
          train(
              agent,
              optimizer,
              processed_gathered_data,
              logger,
              global_step,
              args,
          )
          if global_rank == 0:
              logger.add_scalar(
                  "Time/step_per_second",
                  int(global_step / (time.time() - start_time)),
                  global_step,
              )
  
      envs.close()
      if global_rank == 0:
          test(agent.module, device, logger, args)
  
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 2 for calling a global function with example usage: main(args) without return types