text: import logging
  import argparse
  import configparser
  import os
  import torch
  import torch.nn as nn
  import torch.optim as optim
  
  from crowd_nav.utils.pretrain import freeze_model, trim_model_dict
  from crowd_nav.utils.dataset import ImitDataset, split_dataset
  from crowd_nav.policy.policy_factory import policy_factory
  from crowd_nav.utils.configure import config_log, config_path
  from crowd_nav.snce.contrastive import SocialNCE
  from crowd_nav.snce.model import ProjHead, SpatialEncoder, EventEncoder
  
  torch.manual_seed(2020)
  
  
  def parse_arguments():
      parser = argparse.ArgumentParser('Parse configuration file')
      parser.add_argument('--policy', type=str, default='sail')
      parser.add_argument('--policy_config', type=str, default='configs/policy.config')
      parser.add_argument('--batch_size', type=int, default=128)
      parser.add_argument('--lr', type=float, default=1e-3)
      parser.add_argument('--contrast_sampling', type=str, default='social')
      parser.add_argument('--contrast_weight', type=float, default=1.0)
      parser.add_argument('--contrast_horizon', type=int, default=4)
      parser.add_argument('--contrast_temperature', type=float, default=0.2)
      parser.add_argument('--contrast_range', type=float, default=2.0)
      parser.add_argument('--contrast_nboundary', type=int, default=0)
      parser.add_argument('--ratio_boundary', type=float, default=0.5)
      parser.add_argument('--percent_label', type=float, default=0.5)
      parser.add_argument('--num_epoch', type=int, default=200)
      parser.add_argument('--scheduler_patience', type=int, default=20)
      parser.add_argument('--save_every', type=int, default=5)
      parser.add_argument('--model_file', type=str, default="")
      parser.add_argument('--output_dir', type=str, default='data/output/imitate')
      parser.add_argument('--memory_dir', type=str, default='data/demonstration')
      parser.add_argument('--freeze', default=False, action='store_true')
      parser.add_argument('--gpu', default=False, action='store_true')
      args = parser.parse_args()
      return args
  
  
  def build_policy(args):
      """
      Build navigation policy
      """
      policy = policy_factory[args.policy]()
      if not policy.trainable:
          raise Exception('Policy has to be trainable')
      if args.policy_config is None:
          raise Exception('Policy config has to be specified for a trainable network')
      policy_config = configparser.RawConfigParser()
      policy_config.read(args.policy_config)
      policy.configure(policy_config)
      return policy
  
  
  def set_loader(args, device):
      """
      Set Data Loader
      """
      demo_file = os.path.join(args.memory_dir, 'data_imit.pt')
      logging.info('Load data from %s', demo_file)
      data_imit = torch.load(demo_file)
  
      dataset_imit = ImitDataset(data_imit, None, device, horizon=args.contrast_horizon)
  
      validation_split = 0.3
      DataLoader = split_dataset(dataset_imit, args.batch_size, args.percent_label, validation_split)
      train_loader = DataLoader.train_dataloader()
      valid_loader = DataLoader.val_dataloader()
      #train_loader, valid_loader = split_dataset(dataset_imit, args.batch_size, args.percent_label, validation_split)
      return train_loader, valid_loader
  
  
  def set_model(args, device):
      """
      Set policy network
      """
      policy = build_policy(args)
      policy.set_device(device)
      policy_net = policy.get_model().to(device)
  
      return policy_net
  
  
  def load_model(policy_net, args, device):
      """
      Load pretrained model
      """
      pretrain = torch.load(args.model_file, map_location=device)
      if 'human_encoder.0.weight' in pretrain.keys():
          info = policy_net.load_state_dict(pretrain, strict=True)
      else:
          trim_state = trim_model_dict(pretrain, 'encoder')
          info = policy_net.human_encoder.load_state_dict(trim_state, strict=True)
      print(info)
      logging.info('Load pretrained model from %s', args.model_file)
  
      if args.freeze:
          freeze_model(policy_net, ['human_encoder'])
  
  
  def train(policy_net, projection_head, encoder_sample, train_loader, criterion, nce, optimizer, args):
      """
      Jointly train the policy net and contrastive encoders
      """
      policy_net.train()
      projection_head.train()
      encoder_sample.train()
      loss_sum_all, loss_sum_task, loss_sum_nce = 0, 0, 0
  
      for robot_states, human_states, action_targets, pos_seeds, neg_seeds in train_loader:
  
          # main task
          outputs, features = policy_net(robot_states, human_states)
          loss_task = criterion(outputs, action_targets)
  
          # contrastive task
          if args.contrast_weight > 0:
              loss_nce = nce.loss(robot_states, human_states, pos_seeds, neg_seeds, features)
              loss = loss_task + loss_nce * args.contrast_weight
              loss_sum_nce += loss_nce.item()
          else:
              loss = loss_task
  
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          loss_sum_all += loss.data.item()
          loss_sum_task += loss_task.item()
  
      num_batch = len(train_loader)
      return loss_sum_all / num_batch, loss_sum_task / num_batch, loss_sum_nce / num_batch
  
  
  def validate(policy_net, projection_head, encoder_sample, valid_loader, criterion, nce, args):
      """
      Evaluate policy net
      """
      policy_net.eval()
      projection_head.eval()
      encoder_sample.eval()
      loss_sum_all, loss_sum_task, loss_sum_nce = 0, 0, 0
  
      with torch.no_grad():
          for robot_states, human_states, action_targets, pos_seeds, neg_seeds in valid_loader:
              outputs, features = policy_net(robot_states, human_states)
              loss_task = criterion(outputs, action_targets)
              if args.contrast_weight > 0:
                  loss_nce = nce.loss(robot_states, human_states, pos_seeds, neg_seeds, features)
                  loss = loss_task + loss_nce * args.contrast_weight
                  loss_sum_nce += loss_nce.item()
              else:
                  loss = loss_task
              loss_sum_all += loss.data.item()
              loss_sum_task += loss_task.item()
  
      num_batch = len(valid_loader)
      return loss_sum_all / num_batch, loss_sum_task / num_batch, loss_sum_nce / num_batch
  
  
  def main():
      args = parse_arguments()
      print(args)
  
      # config
      if args.contrast_weight > 0:
          suffix = "-{}-data-{:.2f}-weight-{:.1f}-horizon-{:d}-temperature-{:.2f}-nboundary-{:d}".format(
              args.contrast_sampling, args.percent_label, args.contrast_weight, args.contrast_horizon, args.contrast_temperature, args.contrast_nboundary)
          if args.contrast_nboundary > 0:
              suffix += "-ratio-{:.2f}".format(args.ratio_boundary)
          if args.contrast_sampling == 'local':
              suffix += "-range-{:.2f}".format(args.contrast_range)
      else:
          suffix = "-vanilla-data-{:.2f}".format(args.percent_label)
      config_path(args, suffix)
      config_log(args)
  
      device = torch.device("cuda" if torch.cuda.is_available() and args.gpu else "cpu")
      logging.info('Using device: %s', device)
  
      # dataset
      train_loader, valid_loader = set_loader(args, device)
  
      # model
      policy_net = set_model(args, device)
      projection_head = ProjHead(feat_dim=64, hidden_dim=16, head_dim=8).to(device)
      if args.contrast_sampling == 'event':
          encoder_sample = EventEncoder(hidden_dim=8, head_dim=8).to(device)
      else:
          encoder_sample = SpatialEncoder(hidden_dim=8, head_dim=8).to(device)
  
      # pretrain
      if os.path.exists(args.model_file):
          load_model(policy_net, args, device)
  
      # optimize
      param = list(policy_net.parameters()) + list(projection_head.parameters()) + list(encoder_sample.parameters())
      optimizer = optim.Adam(param, lr=args.lr)
      scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=args.scheduler_patience, threshold=0.01,
                                                       factor=0.5, cooldown=args.scheduler_patience, min_lr=1e-5, verbose=True)
      criterion = nn.MSELoss()
      nce = SocialNCE(projection_head, encoder_sample, args.contrast_sampling, args.contrast_horizon,
                      args.contrast_nboundary, args.contrast_temperature, args.contrast_range, args.ratio_boundary)
  
      # loop
      for epoch in range(args.num_epoch):
  
          train_loss_all, train_loss_task, train_loss_nce = train(policy_net, projection_head, encoder_sample, train_loader, criterion, nce, optimizer, args)
          eval_loss_all, eval_loss_task, eval_loss_nce = validate(policy_net, projection_head, encoder_sample, valid_loader, criterion, nce, args)
  
          scheduler.step(train_loss_all)      # (optional) learning rate decay once training stagnates
  
          if epoch % args.save_every == (args.save_every - 1):
              logging.info("Epoch #%02d: loss = (%.4f, %.4f), task = (%.4f, %.4f), nce = (%.4f, %.4f)", epoch,
                           train_loss_all, eval_loss_all, train_loss_task, eval_loss_task, train_loss_nce, eval_loss_nce)
              torch.save(policy_net.state_dict(), os.path.join(args.output_dir, 'policy_net_{:02d}.pth'.format(epoch)))
  
      torch.save(policy_net.state_dict(), os.path.join(args.output_dir, 'policy_net.pth'))
  
  
  if __name__ == '__main__':
      main()
  