code: def train(total_steps=TOTAL_STEPS):
      # EnvPool runs a batch of environments in separate processes,
      # create_env is a user-defined function that returns a gym environment.
      envs = moolib.EnvPool(
          create_env,
          num_processes=1,
          batch_size=BATCH_SIZE,
          num_batches=1,
      )
  
      model = Model()
  
      broker = moolib.Broker()
      broker.listen(ADDRESS)
      accumulator = moolib.Accumulator(
          "foo", model.parameters(), model.buffers()
      )
      accumulator.connect(ADDRESS)
  
      # Our optimizer.
      opt = torch.optim.Adam(
          model.parameters(), lr=LR, betas=ADAM_BETAS, eps=ADAM_EPSILON
      )
  
      states = []
      actions = []
      rewards = []
      dones = []
      episode_returns = []
  
      action_t = torch.zeros(BATCH_SIZE, dtype=torch.int64)
      episode_step_t = torch.zeros(BATCH_SIZE, dtype=torch.int64)
      episode_return_t = torch.zeros(BATCH_SIZE)
  
      core_state_t = model.initial_state(batch_size=BATCH_SIZE)
      core_states = []
  
      local_loss_computes = 0
      local_update_steps = 0
  
      while True:
          # update does some internal moolib book-keeping,
          # makes sure we're still connected, etc.
          broker.update()
          accumulator.update()
  
          if not accumulator.connected():
              # If we're not connected, sleep for a bit so we don't busy-wait
              print("Your training will commence shortly.")
              time.sleep(0.25)
              continue
  
          # TODO: What's the 0? (Subgroup?)
          obs_future = envs.step(0, action_t)
          obs = obs_future.result()
  
          episode_step_t += 1
          episode_return_t += obs["reward"]
  
          core_states.append(core_state_t)
  
          state_t = obs["state"].to(torch.float, copy=True)
          with torch.no_grad():
              (action_t, logits_t, _), core_state_t = model(
                  state_t, obs["done"], core_state_t
              )
  
          episode_returns.append(episode_return_t.clone())
          episode_step_t *= ~obs["done"]
          episode_return_t *= ~obs["done"]
  
          states.append(state_t)
          actions.append(action_t)
          rewards.append(obs["reward"].clone())
          dones.append(obs["done"].clone())
  
          if accumulator.wants_state():
              # For sharing the initial optimizer state with new nodes.
              accumulator.set_state({"optimizer": opt.state_dict()})
  
          if accumulator.has_new_state():
              # This node wants an initial optimizer state.
              opt.load_state_dict(accumulator.state()["optimizer"])
  
          if accumulator.wants_gradients():
              if len(states) < ROLLOUT_LENGTH + 1:
                  # No data to train on, so this node won't participate in the
                  # gradient reduction this time. Gradients may still be reduced
                  # from the other nodes.
                  accumulator.skip_gradients()
              else:
                  # Train step: We have data to train on. Calculate
                  # a local gradient and initiate a reduction.
                  state = torch.stack(states)
                  action = torch.stack(actions)
                  reward = torch.stack(rewards)
                  done = torch.stack(dones)
  
                  pg_loss, baseline_loss, entropy_loss = a2c_loss(
                      state, action, reward, done, core_states[0], model
                  )
  
                  mean_episode_return = float(
                      "nan"
                  )  # Mean of empty set.
                  if done.any():
                      returns = torch.stack(episode_returns)
                      mean_episode_return = torch.mean(
                          returns[done]
                      ).item()
  
                  del states[:-1]
                  del actions[:-1]
                  del rewards[:-1]
                  del dones[:-1]
                  del episode_returns[:-1]
                  del core_states[:-1]
  
                  loss = (
                      pg_loss
                      + BASELINE_COST * baseline_loss
                      + ENTROPY_COST * entropy_loss
                  )
                  loss.backward()
                  local_loss_computes += 1
  
                  agent_steps = (
                      local_loss_computes * ROLLOUT_LENGTH * BATCH_SIZE
                  )
  
                  if agent_steps >= total_steps:
                      print(
                          "Training ended after %i steps." % agent_steps
                      )
                      return
  
                  grad_norm = torch.nn.utils.clip_grad_norm_(
                      model.parameters(), 100
                  )
  
                  if local_loss_computes % 100 == 0:
                      mer = mean_episode_return
                      print(
                          "Local loss computation %d. (%d agent steps) "
                          "Mean episode return: %s (%d episode ends). "
                          "Local update steps: %d"
                          % (
                              local_loss_computes,
                              agent_steps,
                              None
                              if mer != mer
                              else "%f" % mer,  # NaN -> None.
                              torch.sum(done),
                              local_update_steps,
                          )
                      )
                  log_to_file(
                      step=agent_steps,
                      mean_episode_return=mean_episode_return,
                      pg_loss=pg_loss.item(),
                      baseline_loss=baseline_loss.item(),
                      entropy_loss=entropy_loss.item(),
                      grad_norm=grad_norm.item(),
                  )
                  # Trigger an asynchronous gradient reduction.
                  # has_gradients() will return true when the reduction is done.
                  # We tell it the batch size so we can have virtual batch sizes
                  # and communicate the number of batches consumed.
                  accumulator.reduce_gradients(BATCH_SIZE)
  
          if accumulator.has_gradients():
              # moolib has reduced gradients, so we can step the optimizer
              opt.step()
              accumulator.zero_gradients()  # has_gradients() will return false after this
              local_update_steps += 1
  
file_path: moolib/py/moolib/examples/a2c.py
imports: import csv
  import os
  import time
  import gym
  import torch
  from torch import nn
  from torch.nn import functional
  import moolib
label_desc: Use image node_id 4 for calling a global function with example usage: train(total_steps) without return types