text: import argparse
  import gym
  import rlbench.gym
  import os
  import sys
  import pickle
  import time
  from torch import nn
  from glob import glob
  
  # We need this below line for Python to actually find the utils directory
  # Otherwise, the module is not in the Python path!
  sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
  
  from utils import *
  from algorithms.core.mlp_policy import Policy
  from algorithms.core.mlp_critic import Value
  from algorithms.core.mlp_policy_disc import DiscretePolicy
  from algorithms.core.mlp_discriminator import Discriminator
  from algorithms.core.ppo import ppo_step
  from algorithms.core.common import estimate_advantages
  from algorithms.core.agent import Agent
  
  parser = argparse.ArgumentParser(description='PyTorch GAIL example')
  parser.add_argument('--env-name', default="hit_ball_with_queue-state-v0", metavar='G',
                      help='name of the environment to run')
  parser.add_argument('--expert-traj-path', metavar='G',
                      help='path of the expert trajectories')
  parser.add_argument('--render', action='store_true', default=False,
                      help='render the environment')
  parser.add_argument('--log-std', type=float, default=-0.0, metavar='G',
                      help='log std for the policy (default: -0.0)')
  parser.add_argument('--gamma', type=float, default=0.99, metavar='G',
                      help='discount factor (default: 0.99)')
  parser.add_argument('--tau', type=float, default=0.95, metavar='G',
                      help='gae (default: 0.95)')
  parser.add_argument('--l2-reg', type=float, default=1e-3, metavar='G',
                      help='l2 regularization regression (default: 1e-3)')
  parser.add_argument('--learning-rate', type=float, default=3e-4, metavar='G',
                      help='gae (default: 3e-4)')
  parser.add_argument('--clip-epsilon', type=float, default=0.2, metavar='N',
                      help='clipping epsilon for PPO')
  parser.add_argument('--num-threads', type=int, default=4, metavar='N',
                      help='number of threads for agent (default: 4)')
  parser.add_argument('--seed', type=int, default=1, metavar='N',
                      help='random seed (default: 1)')
  parser.add_argument('--min-batch-size', type=int, default=2048, metavar='N',
                      help='minimal batch size per PPO update (default: 2048)')
  parser.add_argument('--max-iter-num', type=int, default=500, metavar='N',
                      help='maximal number of main iterations (default: 500)')
  parser.add_argument('--log-interval', type=int, default=1, metavar='N',
                      help='interval between training status logs (default: 10)')
  parser.add_argument('--save-model-interval', type=int, default=0, metavar='N',
                      help="interval between saving model (default: 0, means don't save)")
  parser.add_argument('--gpu-index', type=int, default=0, metavar='N')
  args = parser.parse_args()
  
  dtype = torch.float64
  torch.set_default_dtype(dtype)
  device = torch.device('cuda', index=args.gpu_index) if torch.cuda.is_available() else torch.device('cpu')
  if torch.cuda.is_available():
      torch.cuda.set_device(args.gpu_index)
  
  """environment"""
  def key(x):
      return int(x[x.rfind('/')+1:x.rfind('.')])
  
  expert_traj = [sorted(glob(paths+'/*'), key=key) for paths in glob(args.expert_traj_path+'/*')]
  expert_traj = [[pickle.load(open(path, mode='rb')) for path in paths] for paths in expert_traj]
  
  state_dim = expert_traj[0][0].get_low_dim_data().shape[0]
  action_dim = np.hstack([expert_traj[0][0].joint_velocities, expert_traj[0][0].gripper_open]).shape[0]
  
  expert_traj = [[[obs.get_low_dim_data(), obs.joint_velocities, obs.gripper_open] for obs in traj] for traj in expert_traj]
  expert_traj = [np.hstack([traj[i][0], traj[i+1][1], traj[i+1][2]]) for traj in expert_traj for i in range(len(traj)-1)]
  expert_traj = np.vstack(expert_traj)
  #expert_traj = torch.from_numpy(expert_traj).type(dtype)
  
  """seeding"""
  np.random.seed(args.seed)
  torch.manual_seed(args.seed)
  
  """define actor and critic"""
  policy_net = Policy(state_dim, action_dim, log_std=args.log_std, activation='relu')
  value_net = Value(state_dim)
  discrim_net = Discriminator(state_dim + action_dim)
  discrim_criterion = nn.BCELoss()
  to_device(device, policy_net, value_net, discrim_net, discrim_criterion)
  
  optimizer_policy = torch.optim.Adam(policy_net.parameters(), lr=args.learning_rate)
  optimizer_value = torch.optim.Adam(value_net.parameters(), lr=args.learning_rate)
  optimizer_discrim = torch.optim.Adam(discrim_net.parameters(), lr=args.learning_rate)
  
  # optimization epoch number and batch size for PPO
  optim_epochs = 10
  optim_batch_size = 64
  
  
  def expert_reward(state, action):
      state_action = tensor(np.hstack([state, action]), dtype=dtype)
      with torch.no_grad():
          return -math.log(discrim_net(state_action)[0].item())
  
  
  """create agent"""
  agent = Agent(args.env_name, policy_net, device, custom_reward=expert_reward,
                render=args.render, num_threads=args.num_threads)
  
  
  def update_params(batch, i_iter):
      states = torch.from_numpy(np.stack(batch.state)).to(dtype).to(device)
      actions = torch.from_numpy(np.stack(batch.action)).to(dtype).to(device)
      rewards = torch.from_numpy(np.stack(batch.reward)).to(dtype).to(device)
      masks = torch.from_numpy(np.stack(batch.mask)).to(dtype).to(device)
      with torch.no_grad():
          values = value_net(states)
          fixed_log_probs = policy_net.get_log_prob(states, actions)
  
      """get advantage estimation from the trajectories"""
      advantages, returns = estimate_advantages(rewards, masks, values, args.gamma, args.tau, device)
  
      """update discriminator"""
      for _ in range(optim_epochs):
          expert_state_actions = torch.from_numpy(expert_traj).to(dtype).to(device)
          g_o = discrim_net(torch.cat([states, actions], 1))
          e_o = discrim_net(expert_state_actions)
          optimizer_discrim.zero_grad()
          discrim_loss = discrim_criterion(g_o, ones((states.shape[0], 1), device=device)) + \
              discrim_criterion(e_o, zeros((expert_traj.shape[0], 1), device=device))
          discrim_loss.backward()
          optimizer_discrim.step()
  
      """perform mini-batch PPO update"""
      optim_iter_num = int(math.ceil(states.shape[0] / optim_batch_size))
      for _ in range(optim_epochs):
          perm = torch.randperm(states.shape[0]).to(device)
  
          states, actions, returns, advantages, fixed_log_probs = \
              states[perm].clone(), actions[perm].clone(), returns[perm].clone(), advantages[perm].clone(), fixed_log_probs[perm].clone()
  
          for i in range(optim_iter_num):
              ind = slice(i * optim_batch_size, min((i + 1) * optim_batch_size, states.shape[0]))
              states_b, actions_b, advantages_b, returns_b, fixed_log_probs_b = \
                  states[ind], actions[ind], advantages[ind], returns[ind], fixed_log_probs[ind]
  
              ppo_step(policy_net, value_net, optimizer_policy, optimizer_value, 1, states_b, actions_b, returns_b,
                       advantages_b, fixed_log_probs_b, args.clip_epsilon, args.l2_reg)
  
  
  def main_loop():
      for i_iter in range(args.max_iter_num):
          """generate multiple trajectories that reach the minimum batch_size"""
          discrim_net.to(torch.device('cpu'))
          batch, log = agent.collect_samples(args.min_batch_size)
          discrim_net.to(device)
  
          t0 = time.time()
          update_params(batch, i_iter)
          t1 = time.time()
  
          if i_iter % args.log_interval == 0:
              print('=======================================')
              print('=======================================')
              print('{}\tT_sample {:.4f}\tT_update {:.4f}\texpert_R_avg {:.2f}\tR_avg {:.2f}'.format(
                  i_iter, log['sample_time'], t1-t0, log['avg_c_reward'], log['avg_reward']))
              print('=======================================')
              print('=======================================')
  
          if args.save_model_interval > 0 and (i_iter+1) % args.save_model_interval == 0:
              to_device(torch.device('cpu'), policy_net, value_net, discrim_net)
              pickle.dump((policy_net, value_net, discrim_net), open(os.path.join(assets_dir(), 'learned_models/{}_gail.p'.format(args.env_name)), 'wb'))
              to_device(device, policy_net, value_net, discrim_net)
  
          """clean up gpu memory"""
          torch.cuda.empty_cache()
  
  
  main_loop()
  