code: def main():
      logger.info("------------------- CQL ---------------------")
      logger.info("Env: {}, Seed: {}".format(args.env, args.seed))
      logger.info("---------------------------------------------")
  
      env = gym.make(args.env)
      env.seed(args.seed)
  
      obs_dim = env.observation_space.shape[0]
      action_dim = env.action_space.shape[0]
  
      # Initialize model, algorithm, agent
      model = MujocoModel(obs_dim, action_dim)
      algorithm = CQL(
          model,
          gamma=GAMMA,
          tau=TAU,
          actor_lr=ACTOR_LR,
          critic_lr=CRITIC_LR,
          with_automatic_entropy_tuning=args.with_automatic_entropy_tuning,
          with_lagrange=args.with_lagrange,
      )
      agent = MujocoAgent(algorithm)
  
      # Initialize offline data
      rpm = ReplayMemory(
          max_size=MEMORY_SIZE, obs_dim=obs_dim, act_dim=action_dim
      )
      rpm.load_from_d4rl(d4rl.qlearning_dataset(env))
  
      for total_steps in trange(int(args.train_total_steps)):
          # Train steps
          (
              batch_obs,
              batch_action,
              batch_reward,
              batch_next_obs,
              batch_terminal,
          ) = rpm.sample_batch(BATCH_SIZE)
          critic_loss, mse_loss, actor_loss, min_q = agent.learn(
              batch_obs,
              batch_action,
              batch_reward,
              batch_next_obs,
              batch_terminal,
          )
          summary.add_scalar(
              "train/critic_loss",
              critic_loss.cpu().numpy(),
              total_steps,
          )
  
          summary.add_scalar(
              "train/mse_loss", mse_loss.cpu().numpy(), total_steps
          )
          summary.add_scalar(
              "train/actor_loss", actor_loss.cpu().numpy(), total_steps
          )
          summary.add_scalar(
              "train/min_qi", min_q.cpu().numpy(), total_steps
          )
  
          # Evaluate episode
          if total_steps % args.test_every_steps == 0:
              avg_reward = run_evaluate_episodes(
                  agent, env, EVAL_EPISODES
              )
              summary.add_scalar(
                  "eval/episode_reward", avg_reward, total_steps
              )
              logger.info(
                  "Evaluation: total_steps {}, Reward: {}".format(
                      total_steps, avg_reward
                  )
              )
  
file_path: PARL/examples/CQL/train.py
imports: import argparse
  import gym
  import d4rl
  from parl.utils import logger, summary, ReplayMemory
  from mujoco_model import MujocoModel
  from mujoco_agent import MujocoAgent
  from parl.algorithms import CQL
  from tqdm import trange
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types