code: def main(config):
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
  
      env_name, dataset = config["env"], config["dataset"]
  
      if env_name == "hopper":
          env = gym.make("Hopper-v3")
          env_targets = 3600
      elif env_name == "halfcheetah":
          env = gym.make("HalfCheetah-v3")
          env_targets = 12000
      elif env_name == "walker2d":
          env = gym.make("Walker2d-v3")
          env_targets = 5000
      else:
          raise NotImplementedError
  
      state_dim = env.observation_space.shape[0]
      act_dim = env.action_space.shape[0]
  
      logger.info(
          "Starting new experiment: {} {}".format(env_name, dataset)
      )
  
      K = config["K"]
      batch_size = config["batch_size"]
      num_eval_episodes = config["num_eval_episodes"]
      pct_traj = config.get("pct_traj", 1.0)
  
      model = TrajectoryModel(
          state_dim=state_dim,
          act_dim=act_dim,
          max_length=K,
          max_ep_len=config["max_ep_len"],
          hidden_size=config["embed_dim"],
          n_layer=config["n_layer"],
          n_head=config["n_head"],
          n_inner=4 * config["embed_dim"],
          activation_function=config["activation_function"],
          n_positions=1024,
          resid_pdrop=config["dropout"],
          attn_pdrop=config["dropout"],
      )
  
      model = model.to(device=device)
      alg = parl.algorithms.DecisionTransformer(
          model,
          config["learning_rate"],
          config["warmup_steps"],
          config["weight_decay"],
      )
      agent = DTAgent(alg, config)
  
      dataset_path = f"data/{env_name}-{dataset}-v2.pkl"
      state_mean, state_std = agent.load_data(dataset_path)
  
      for iter in range(config["max_iters"]):
          train_loss = []
          agent.train()
          for step in tqdm(range(config["num_steps_per_iter"])):
              loss = agent.learn()
              train_loss.append(loss)
          logger.info(
              "[training] iter:{} loss:{}".format(
                  iter, np.mean(train_loss)
              )
          )
          summary.add_scalar("train_loss", np.mean(train_loss), iter)
  
          agent.eval()
          logs = eval_episodes(
              env_targets,
              env,
              state_dim,
              act_dim,
              agent,
              config["max_ep_len"],
              config["rew_scale"],
              state_mean,
              state_std,
              device,
          )
          for k, v in logs.items():
              logger.info("{}: {}".format(k, v))
              summary.add_scalar(k, v, iter)
  
file_path: PARL/benchmark/torch/DT/train.py
imports: import gym
  import numpy
  import torch
  import argparse
  import pickle
  import random
  import sys
  from tqdm import tqdm
  import parl
  from model import TrajectoryModel
  from agent import DTAgent
  from parl.utils import summary
  from evaluate_episodes import eval_episodes
  from parl.utils import logger
label_desc: Use image node_id 1 for calling a global function with example usage: main(config) without return types