code: def main():
      config = mujoco_config if args.continuous_action else atari_config
      if args.env_num:
          config["env_num"] = args.env_num
      config["env"] = args.env
      config["seed"] = args.seed
      config["xparl_addr"] = args.xparl_addr
      config["test_every_steps"] = args.test_every_steps
  
      config["batch_size"] = int(
          config["env_num"] * config["step_nums"]
      )
      config["num_updates"] = int(
          config["train_total_steps"] // config["batch_size"]
      )
  
      logger.info("------------------- PPO ---------------------")
      logger.info(
          "Env: {}, seed: {}".format(config["env"], config["seed"])
      )
      logger.info("---------------------------------------------")
      logger.set_dir(
          "./train_logs/{}_{}".format(config["env"], config["seed"])
      )
  
      envs = ParallelEnv(config)
      eval_env = LocalEnv(config["env"], test=True)
  
      obs_space = eval_env.obs_space
      act_space = eval_env.act_space
  
      if config["continuous_action"]:
          model = MujocoModel(obs_space, act_space)
      else:
          model = AtariModel(obs_space, act_space)
      ppo = PPO(
          model,
          clip_param=config["clip_param"],
          entropy_coef=config["entropy_coef"],
          initial_lr=config["initial_lr"],
          continuous_action=config["continuous_action"],
      )
      agent = PPOAgent(ppo, config)
  
      rollout = RolloutStorage(
          config["step_nums"], config["env_num"], obs_space, act_space
      )
  
      obs = envs.reset()
      done = np.zeros(config["env_num"], dtype="float32")
  
      test_flag = 0
      total_steps = 0
      for update in range(1, config["num_updates"] + 1):
          for step in range(0, config["step_nums"]):
              total_steps += 1 * config["env_num"]
  
              value, action, logprob, _ = agent.sample(obs)
              next_obs, reward, next_done, info = envs.step(action)
              rollout.append(
                  obs, action, logprob, reward, done, value.flatten()
              )
              obs, done = next_obs, next_done
  
              for k in range(config["env_num"]):
                  if done[k] and "episode" in info[k].keys():
                      logger.info(
                          "Training: total steps: {}, episode rewards: {}".format(
                              total_steps, info[k]["episode"]["r"]
                          )
                      )
                      summary.add_scalar(
                          "train/episode_reward",
                          info[k]["episode"]["r"],
                          total_steps,
                      )
  
          # Bootstrap value if not done
          value = agent.value(obs)
          rollout.compute_returns(value, done)
  
          # Optimizing the policy and value network
          v_loss, pg_loss, entropy_loss, lr = agent.learn(rollout)
  
          if total_steps // config["test_every_steps"] >= test_flag:
              while (
                  total_steps // config["test_every_steps"] >= test_flag
              ):
                  test_flag += 1
  
              if config["continuous_action"]:
                  # set running mean and variance of obs
                  ob_rms = envs.eval_ob_rms
                  eval_env.env.set_ob_rms(ob_rms)
  
              avg_reward = run_evaluate_episodes(
                  agent, eval_env, config["eval_episode"]
              )
              summary.add_scalar(
                  "eval/episode_reward", avg_reward, total_steps
              )
              logger.info(
                  "Evaluation over: {} episodes, Reward: {}".format(
                      config["eval_episode"], avg_reward
                  )
              )
  
file_path: PARL/benchmark/torch/ppo/train.py
imports: import argparse
  import numpy
  from parl.utils import logger, summary
  from mujoco_config import mujoco_config
  from atari_config import atari_config
  from env_utils import ParallelEnv, LocalEnv
  from storage import RolloutStorage
  from atari_model import AtariModel
  from mujoco_model import MujocoModel
  from parl.algorithms import PPO
  from agent import PPOAgent
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types