content: import gc
  import os
  import io
  import sys
  import random
  import warnings
  from collections import defaultdict
  from typing import Dict, List
  import jsonlines
  
  import lmdb
  import msgpack_numpy
  import numpy as np
  import math
  import time
  import torch
  import torch.nn.functional as F
  from torch.autograd import Variable
  from torch.nn.parallel import DistributedDataParallel as DDP
  
  import tqdm
  from gym import Space
  from habitat import Config, logger
  from habitat_baselines.common.baseline_registry import baseline_registry
  from habitat_baselines.common.environments import get_env_class
  from habitat_baselines.common.obs_transformers import (
      apply_obs_transforms_batch,
      apply_obs_transforms_obs_space,
      get_active_obs_transforms,
  )
  from habitat_baselines.common.tensorboard_utils import TensorboardWriter
  from habitat_baselines.utils.common import batch_obs
  
  from vlnce_baselines.common.aux_losses import AuxLosses
  from vlnce_baselines.common.base_il_trainer import BaseVLNCETrainer
  from vlnce_baselines.common.env_utils import construct_envs, construct_envs_for_rl, is_slurm_batch_job
  from vlnce_baselines.common.utils import extract_instruction_tokens
  from vlnce_baselines.utils import reduce_loss
  
  from habitat.utils.visualizations.utils import images_to_video
  
  from .utils import get_camera_orientations12
  from .utils import (
      length2mask, dir_angle_feature_with_ele,
  )
  from vlnce_baselines.common.utils import dis_to_con, gather_list_and_concat
  from habitat_extensions.measures import NDTW
  from fastdtw import fastdtw
  
  with warnings.catch_warnings():
      warnings.filterwarnings("ignore", category=FutureWarning)
      import tensorflow as tf  # noqa: F401
  
  import torch.distributed as distr
  import gzip
  import json
  from copy import deepcopy
  from torch.cuda.amp import autocast, GradScaler
  
  @baseline_registry.register_trainer(name="HAMT")
  class RLTrainer(BaseVLNCETrainer):
      def __init__(self, config=None):
          super().__init__(config)
          self.max_len = int(config.IL.max_traj_len) #  * 0.97 transfered gt path got 0.96 spl
  
          #check ceph things
          if config.CEPH_IO:
              from petrel_client.client import Client
              conf_path = '~/petreloss.conf'
              self.client = Client(conf_path)
  
      def _make_dirs(self) -> None:
          if self.config.local_rank == 0:
              self._make_ckpt_dir()
              # os.makedirs(self.lmdb_features_dir, exist_ok=True)
              if self.config.EVAL.SAVE_RESULTS:
                  self._make_results_dir()
  
      def save_checkpoint(self, iteration: int,) -> None:
          if not self.config.CEPH_IO:
              torch.save(
                  obj={
                      "state_dict": self.policy.state_dict(),
                      "config": self.config,
                      "optim_state": self.optimizer.state_dict(),
                      "iteration": iteration,
                  },
                  f=os.path.join(self.config.CHECKPOINT_FOLDER, f"ckpt.iter{iteration}.pth"),
              )
          else:
              save_dict = {
                      "state_dict": self.policy.state_dict(),
                      "config": self.config,
                      "optim_state": self.optimizer.state_dict(),
                      "iteration": iteration,
                      }
              path = os.path.join(self.config.CEPH_URL, f"ckpt.iter{iteration}.pth")
              with io.BytesIO() as buffer:
                  torch.save(save_dict, buffer)
                  self.client.put(path, buffer.getvalue())
  
      def _set_config(self):
          self.split = self.config.TASK_CONFIG.DATASET.SPLIT
          self.config.defrost()
          self.config.TASK_CONFIG.TASK.NDTW.SPLIT = self.split
          self.config.TASK_CONFIG.TASK.SDTW.SPLIT = self.split
          self.config.TASK_CONFIG.TASK.MEASUREMENTS = ['DISTANCE_TO_GOAL', 'NDTW'] # for RL reward
          self.config.TASK_CONFIG.ENVIRONMENT.MAX_EPISODE_STEPS = self.config.IL.max_traj_len 
          self.config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_STEPS = -1
          self.config.SIMULATOR_GPU_IDS = self.config.SIMULATOR_GPU_IDS[self.config.local_rank]
          self.config.use_pbar = True # not is_slurm_batch_job()
          ''' if choosing image '''
          resize_config = self.config.RL.POLICY.OBS_TRANSFORMS.RESIZER_PER_SENSOR.SIZES
          crop_config = self.config.RL.POLICY.OBS_TRANSFORMS.CENTER_CROPPER_PER_SENSOR.SENSOR_CROPS
          config = self.config.TASK_CONFIG
          camera_orientations = get_camera_orientations12()
          for sensor_type in ["RGB", "DEPTH"]:
              resizer_size = dict(resize_config)[sensor_type.lower()]
              cropper_size = dict(crop_config)[sensor_type.lower()]
              sensor = getattr(config.SIMULATOR, f"{sensor_type}_SENSOR")
              for action, orient in camera_orientations.items():
                  camera_template = f"{sensor_type}_{action}"
                  camera_config = deepcopy(sensor)
                  camera_config.ORIENTATION = camera_orientations[action]
                  camera_config.UUID = camera_template.lower()
                  setattr(config.SIMULATOR, camera_template, camera_config)
                  config.SIMULATOR.AGENT_0.SENSORS.append(camera_template)
                  resize_config.append((camera_template.lower(), resizer_size))
                  crop_config.append((camera_template.lower(), cropper_size))
          self.config.RL.POLICY.OBS_TRANSFORMS.RESIZER_PER_SENSOR.SIZES = resize_config
          self.config.RL.POLICY.OBS_TRANSFORMS.CENTER_CROPPER_PER_SENSOR.SENSOR_CROPS = crop_config
          self.config.TASK_CONFIG = config
          self.config.SENSORS = config.SIMULATOR.AGENT_0.SENSORS
  
          if self.config.IL.progress_monitor == True:
              self.config.MODEL.progress_monitor = True
              self.config.MODEL.max_len = self.config.IL.max_text_len
          else:
              self.config.MODEL.progress_monitor = False
  
          self.config.freeze()
  
  
  
          self.world_size = self.config.GPU_NUMBERS
          self.local_rank = self.config.local_rank
          self.batch_size = self.config.IL.batch_size
          torch.cuda.set_device(self.device)
          if self.world_size > 1:
              distr.init_process_group(backend='nccl', init_method='env://')
              self.device = self.config.TORCH_GPU_IDS[self.local_rank]
              self.config.defrost()
              self.config.TORCH_GPU_ID = self.config.TORCH_GPU_IDS[self.local_rank]
              self.config.freeze()
              torch.cuda.set_device(self.device)
  
      def _init_envs(self):
          self.config.defrost()
          self.config.TASK_CONFIG.SEED = self.config.TASK_CONFIG.SEED + self.local_rank
          self.config.freeze()
  
          self.envs = construct_envs(
              self.config, 
              get_env_class(self.config.ENV_NAME),
              auto_reset_done=False
          )
          env_num = self.envs.num_envs
          dataset_len = sum(self.envs.number_of_episodes)
          logger.info(f'LOCAL RANK: {self.local_rank}, ENV NUM: {env_num}, DATASET LEN: {dataset_len}')
          observation_space = self.envs.observation_spaces[0]
          action_space = self.envs.action_spaces[0]
          self.obs_transforms = get_active_obs_transforms(self.config)
          observation_space = apply_obs_transforms_obs_space(
              observation_space, self.obs_transforms
          )
  
          return observation_space, action_space
  
      def _get_iter(self, x):
          x_iter = int(x.split('.')[-2][4:])
          return x_iter
  
      def _initialize_policy(
          self,
          config: Config,
          load_from_ckpt: bool,
          observation_space: Space,
          action_space: Space,
      ) -> int:
          start_iter = 0
          policy = baseline_registry.get_policy(self.config.MODEL.policy_name)
          self.policy = policy.from_config(
              config=config,
              observation_space=observation_space,
              action_space=action_space,
          )
          ''' initialize the waypoint predictor here '''
          from vlnce_baselines.waypoint_pred.TRM_net import BinaryDistPredictor_TRM
          self.waypoint_predictor = BinaryDistPredictor_TRM(device=self.device)
          self.waypoint_predictor.load_state_dict(
              torch.load('pretrained/wp_pred/waypoint_predictor', map_location = torch.device('cpu'))['predictor']['state_dict']
          )
          for param in self.waypoint_predictor.parameters():
              param.requires_grad_(False)
  
          self.policy.to(self.device)
          self.waypoint_predictor.to(self.device)
          self.num_recurrent_layers = self.policy.net.num_recurrent_layers
  
          if self.config.GPU_NUMBERS > 1:
              print('Using', self.config.GPU_NUMBERS,'GPU!')
              # find_unused_parameters=False fix ddp bug
              self.policy.net = DDP(self.policy.net.to(self.device), device_ids=[self.device],
                  output_device=self.device, find_unused_parameters=False, broadcast_buffers=False)
          self.optimizer = torch.optim.AdamW(self.policy.parameters(), lr=self.config.IL.lr)
  
          if config.IL.resume:
              import glob
              if not self.config.CEPH_IO:
                  ckpt_list = list(filter(os.path.isfile, glob.glob(config.CHECKPOINT_FOLDER + "/*")) )
              else:
                  ckpt_list = [os.path.join(self.config.CEPH_URL,p) for p in self.client.list(self.config.CEPH_URL)]
  
              ckpt_list.sort(key=self._get_iter)
  
              if len(ckpt_list) > 0:
                  config.defrost()
                  config.IL.ckpt_to_load = ckpt_list[-1]
                  load_from_ckpt = True
                  config.IL.is_requeue = True
                  config.freeze()
              else:
                  load_from_ckpt = False
  
          if load_from_ckpt:
              ckpt_path = config.IL.ckpt_to_load
  
              if not self.config.CEPH_IO:
                  ckpt_dict = self.load_checkpoint(ckpt_path, map_location="cpu")
              else:
                  with io.BytesIO(self.client.get(ckpt_path)) as buffer:
                      ckpt_dict = torch.load(buffer, map_location="cpu")
  
              if 'module' in list(ckpt_dict['state_dict'].keys())[0] and self.config.GPU_NUMBERS == 1:
                  self.policy.net = torch.nn.DataParallel(self.policy.net.to(self.device),
                      device_ids=[self.device], output_device=self.device)
                  self.policy.load_state_dict(ckpt_dict["state_dict"], strict=False)
                  self.policy.net = self.policy.net.module
                  self.waypoint_predictor = torch.nn.DataParallel(self.waypoint_predictor.to(self.device),
                      device_ids=[self.device], output_device=self.device)
              else:
                  self.policy.load_state_dict(ckpt_dict["state_dict"], strict=False)
              if config.IL.is_requeue:
                  start_iter = ckpt_dict["iteration"]
                  self.optimizer.load_state_dict(ckpt_dict["optim_state"])
              logger.info(f"Loaded weights from checkpoint: {ckpt_path}, iteration: {start_iter}")
  		
          params = sum(param.numel() for param in self.policy.parameters())
          params_t = sum(
              p.numel() for p in self.policy.parameters() if p.requires_grad
          )
          logger.info(f"Agent parameters: {params/1e6:.2f} MB. Trainable: {params_t/1e6:.2f} MB.")
          logger.info("Finished setting up policy.")
  
          return start_iter
  
      def _teacher_action(self, batch_angles, batch_distances, candidate_lengths):
          if self.config.MODEL.task_type == 'r2r':
              cand_dists_to_goal = [[] for _ in range(len(batch_angles))]
              oracle_cand_idx = []
              for j in range(len(batch_angles)):
                  for k in range(len(batch_angles[j])):
                      angle_k = batch_angles[j][k]
                      forward_k = batch_distances[j][k]
                      dist_k = self.envs.call_at(j, "cand_dist_to_goal", {"angle": angle_k, "forward": forward_k})
                      cand_dists_to_goal[j].append(dist_k)
                  curr_dist_to_goal = self.envs.call_at(j, "current_dist_to_goal")
                  # if within target range (which def as 3.0)
                  if curr_dist_to_goal < 1.5:
                      oracle_cand_idx.append(candidate_lengths[j] - 1)
                  else:
                      oracle_cand_idx.append(np.argmin(cand_dists_to_goal[j]))
              return oracle_cand_idx
          elif self.config.MODEL.task_type == 'rxr':
              kargs = []
              current_episodes = self.envs.current_episodes()
              for i in range(self.envs.num_envs):
                  kargs.append({
                      'ref_path':self.gt_data[str(current_episodes[i].episode_id)]['locations'],
                      'angles':batch_angles[i],
                      'distances':batch_distances[i],
                      'candidate_length':candidate_lengths[i]
                  })
  
              outputs = self.envs.call(["get_cand_idx"]*self.envs.num_envs,kargs)
              oracle_cand_idx, progresses = [list(x) for x in zip(*outputs)]
              return oracle_cand_idx, progresses
  
      def _cand_pano_feature_variable(self, obs):
          batch_size = len(obs['cand_angles'])
          ob_cand_lens = [len(x)+1 for x in obs['cand_angles']]  # +1 is for the end
          ob_lens = []
          ob_rgb_fts, ob_dep_fts, ob_ang_fts, ob_dis_fts, ob_nav_types = [], [], [], [], []
          # Note: The candidate_feat at len(ob['candidate']) is the feature for the END
          # which is zero in my implementation
          for i in range(batch_size):
              cand_nav_types = []
              cand_idxes = np.zeros(12, dtype=np.bool)
              cand_idxes[obs['cand_img_idxes'][i]] = True
              # cand
              cand_rgb_fts = obs['cand_rgb'][i]
              cand_dep_fts = obs['cand_depth'][i]
              cand_ang_fts = obs['cand_angle_fts'][i]
              cand_dis_fts = obs['cand_dis_fts'][i]
              cand_nav_types += [1] * cand_ang_fts.shape[0]
              # stop
              stop_rgb_fts = torch.zeros([1, 768])
              # stop_rgb_fts = torch.zeros([1, 2048])
              stop_dep_fts = torch.zeros([1, 128])
              stop_ang_fts = torch.zeros([1, 4])
              stop_dis_fts = torch.zeros([1, 4])
              cand_nav_types += [2]
              # pano context
              pano_rgb_fts = obs['pano_rgb'][i][~cand_idxes]
              pano_dep_fts = obs['pano_depth'][i][~cand_idxes]
              pano_ang_fts = obs['pano_angle_fts'][~cand_idxes]
              pano_dis_fts = obs['pano_dis_fts'][~cand_idxes]
              cand_nav_types += [0] * (12-np.sum(cand_idxes))
  
              cand_pano_rgb = torch.cat([cand_rgb_fts, stop_rgb_fts, pano_rgb_fts], dim=0)
              cand_pano_dep = torch.cat([cand_dep_fts, stop_dep_fts, pano_dep_fts], dim=0)
              cand_pano_ang = torch.cat([cand_ang_fts, stop_ang_fts, pano_ang_fts], dim=0)
              cand_pano_dis = torch.cat([cand_dis_fts, stop_dis_fts, pano_dis_fts], dim=0)
              ob_rgb_fts.append(cand_pano_rgb)
              ob_dep_fts.append(cand_pano_dep)
              ob_ang_fts.append(cand_pano_ang)
              ob_dis_fts.append(cand_pano_dis)
              ob_nav_types.append(cand_nav_types)
              ob_lens.append(len(cand_nav_types))
          
          # pad features to max_len
          max_len = max(ob_lens)
          for i in range(batch_size):
              num_pads = max_len - ob_lens[i]
              ob_rgb_fts[i] = torch.cat([ob_rgb_fts[i], torch.zeros(num_pads, 768)], dim=0)
              # ob_rgb_fts[i] = torch.cat([ob_rgb_fts[i], torch.zeros(num_pads, 2048)], dim=0)
              ob_dep_fts[i] = torch.cat([ob_dep_fts[i], torch.zeros(num_pads, 128)], dim=0)
              ob_ang_fts[i] = torch.cat([ob_ang_fts[i], torch.zeros(num_pads, 4)], dim=0)
              ob_dis_fts[i] = torch.cat([ob_dis_fts[i], torch.zeros(num_pads, 4)], dim=0)
              ob_nav_types[i] = np.array(ob_nav_types[i] + [0]*num_pads)
          
          ob_rgb_fts = torch.stack(ob_rgb_fts, dim=0).cuda()
          ob_dep_fts = torch.stack(ob_dep_fts, dim=0).cuda()
          ob_ang_fts = torch.stack(ob_ang_fts, dim=0).cuda()
          ob_dis_fts = torch.stack(ob_dis_fts, dim=0).cuda()
          ob_nav_types = torch.from_numpy(np.stack(ob_nav_types, 0)).cuda()
  
          return ob_rgb_fts, ob_dep_fts, ob_ang_fts, ob_dis_fts, ob_nav_types, ob_lens, ob_cand_lens
  
      def _history_variable(self, obs):
          batch_size = obs['pano_rgb'].shape[0]
          hist_rgb_fts = obs['pano_rgb'][:, 0, ...].cuda()
          hist_depth_fts = obs['pano_depth'][:, 0, ...].cuda()
          hist_pano_rgb_fts = obs['pano_rgb'].cuda()
          hist_pano_depth_fts = obs['pano_depth'].cuda()
          hist_pano_ang_fts = obs['pano_angle_fts'].unsqueeze(0).expand(batch_size, -1, -1).cuda()
  
          return hist_rgb_fts, hist_depth_fts, hist_pano_rgb_fts, hist_pano_depth_fts, hist_pano_ang_fts
  
      def _train_iter(self, train_ml=None, sample_ratio=None, train_rl=True):
          if train_ml is not None:
              feedback = 'teacher'
              assert sample_ratio is not None
          elif train_rl:
              feedback = 'sample'
              assert sample_ratio is None
          else:
              raise NotImplementedError
  
          
          self.envs.resume_all()
          observations = self.envs.reset()
          instr_max_len = self.config.IL.max_text_len # r2r 80, rxr 200
          instr_pad_id = 1 if self.config.MODEL.task_type == 'rxr' else 0
          observations = extract_instruction_tokens(observations, self.config.TASK_CONFIG.TASK.INSTRUCTION_SENSOR_UUID,
                                                    max_length=instr_max_len, pad_id=instr_pad_id)
          batch = batch_obs(observations, self.device)
          batch = apply_obs_transforms_batch(batch, self.obs_transforms)
          keys = ['rgb', 'rgb_30', 'rgb_60', 'rgb_90', 'rgb_120', 'rgb_150', 'rgb_180', 'rgb_210', 'rgb_240', 'rgb_270', 'rgb_300', 'rgb_330']
          # states = [[self.envs.call_at(i,"get_agent_state",{})] for i in range(self.envs.num_envs)]
          history_images = [{k:observations[i][k][None,...] for k in keys} for i in range(self.envs.num_envs)]
          video_inputs = [{k:observations[i][k][None,...].repeat(16,0) for k in keys} for i in range(self.envs.num_envs)]
          batch['video_rgbs'] = video_inputs
  
  
          # encode instructions
          all_txt_ids = batch['instruction']
          all_txt_masks = (all_txt_ids != instr_pad_id)
          all_txt_embeds = self.policy.net(
              mode='language',
              txt_ids=all_txt_ids,
              txt_masks=all_txt_masks,
          )
  
          il_loss = 0.
          total_actions = 0.
          not_done_index = list(range(self.envs.num_envs))
          
          # for rl
          last_dist = np.zeros(self.envs.num_envs, np.float32)
          last_ndtw = np.zeros(self.envs.num_envs, np.float32)
          for i in range(self.envs.num_envs):
              info = self.envs.call_at(i, "get_metrics", {})
              last_dist[i] = info['distance_to_goal']
              last_ndtw[i] = info['ndtw']
          rewards = []
          entropys = []
          hidden_states = []
          policy_log_probs = []
          not_done_index_bank = []
          init_num_envs = self.envs.num_envs
  
          hist_lens = np.ones(self.envs.num_envs, dtype=np.int64)
          hist_embeds = [self.policy.net('history').expand(self.envs.num_envs, -1)]
          
          for stepk in range(self.max_len):
              total_actions += self.envs.num_envs
              txt_embeds = all_txt_embeds[not_done_index]
              txt_masks = all_txt_masks[not_done_index]
              
              # cand waypoint prediction
              wp_outputs = self.policy.net(
                  mode = "waypoint",
                  waypoint_predictor = self.waypoint_predictor,
                  observations = batch,
                  in_train = self.config.IL.waypoint_aug,
              )
              ob_rgb_fts, ob_dep_fts, ob_ang_fts, ob_dis_fts, \
              ob_nav_types, ob_lens, ob_cand_lens = self._cand_pano_feature_variable(wp_outputs)
              ob_masks = length2mask(ob_lens).logical_not()
  
              # navigation
              visual_inputs = {
                  'mode': 'navigation',
                  'txt_embeds': txt_embeds,
                  'txt_masks': txt_masks,
                  'hist_embeds': hist_embeds,    # history before t step
                  'hist_lens': hist_lens,
                  'ob_rgb_fts': ob_rgb_fts,
                  'ob_dep_fts': ob_dep_fts,
                  'ob_ang_fts': ob_ang_fts,
                  'ob_dis_fts': ob_dis_fts,
                  'ob_nav_types': ob_nav_types,
                  'ob_masks': ob_masks,
                  'return_states': True if feedback == 'sample' or self.config.IL.progress_monitor else False,
              }
              with autocast():
                  t_outputs = self.policy.net(**visual_inputs)
                  logits = t_outputs[0]
  
                  if feedback == 'sample':
                      h_t = t_outputs[1]
                      hidden_states.append(h_t)
  
                  if train_ml is not None:
                      if self.config.MODEL.task_type == 'r2r':
                          oracle_cand_idx = self._teacher_action(wp_outputs['cand_angles'], wp_outputs['cand_distances'], ob_cand_lens)
                      elif self.config.MODEL.task_type == 'rxr':
                          oracle_cand_idx, gt_pm = self._teacher_action(wp_outputs['cand_angles'], wp_outputs['cand_distances'], ob_cand_lens)
                      oracle_actions = torch.tensor(oracle_cand_idx, device=self.device)
                      il_loss += F.cross_entropy(logits, oracle_actions, reduction='sum')
                      if self.config.IL.progress_monitor:
                          h_t = t_outputs[1]
                          language_attention =  t_outputs[2]
                          progresses = self.policy.net(mode='progress',h_t=h_t,language_attention=language_attention)
                          gt_pm = torch.tensor(gt_pm,device=self.device,dtype=torch.float32).unsqueeze(1)
                          pm_loss = F.mse_loss(progresses, gt_pm)
                          il_loss += torch.sum(pm_loss)
  
              # determine action
              if feedback == 'teacher':
                  # a_t = oracle_actions
                  a_t = logits.argmax(dim=-1)
                  a_t = torch.where(torch.rand_like(a_t, dtype=torch.float)<=sample_ratio, oracle_actions, a_t)
              elif feedback == 'sample':
                  probs = F.softmax(logits, 1)  # sampling an action from model
                  c = torch.distributions.Categorical(probs)
                  self.logs['entropy'].append(c.entropy().mean().item())            # For log
                  entropys.append(c.entropy())                                     # For optimization
                  a_t = c.sample().detach()
                  policy_log_probs.append(c.log_prob(a_t))
              else:
                  raise NotImplementedError
              cpu_a_t = a_t.cpu().numpy()
  
              # update history
              if stepk != self.max_len-1:
                  hist_rgb_fts, hist_depth_fts, hist_pano_rgb_fts, hist_pano_depth_fts, hist_pano_ang_fts = self._history_variable(wp_outputs)
                  prev_act_ang_fts = torch.zeros([self.envs.num_envs, 4]).cuda()
                  for i, next_id in enumerate(cpu_a_t):
                      prev_act_ang_fts[i] = ob_ang_fts[i, next_id]
                  t_hist_inputs = {
                      'mode': 'history',
                      'hist_rgb_fts': hist_rgb_fts,
                      'hist_depth_fts': hist_depth_fts,
                      'hist_ang_fts': prev_act_ang_fts,
                      'hist_pano_rgb_fts': hist_pano_rgb_fts,
                      'hist_pano_depth_fts': hist_pano_depth_fts,
                      'hist_pano_ang_fts': hist_pano_ang_fts,
                      'ob_step': stepk,
                  }
                  with autocast():
                      t_hist_embeds = self.policy.net(**t_hist_inputs)
                      hist_embeds.append(t_hist_embeds)
                  hist_lens = hist_lens + 1
  
              # make equiv action
              env_actions = []
              for j in range(self.envs.num_envs):
                  if cpu_a_t[j].item()==ob_cand_lens[j]-1 or stepk==self.max_len-1:
                      env_actions.append({'action':{'action': 0, 'action_args':{}}})
                  else:
                      t_angle = wp_outputs['cand_angles'][j][cpu_a_t[j]]
                      t_distance = wp_outputs['cand_distances'][j][cpu_a_t[j]]
                      env_actions.append({'action':{'action': 4, 'action_args':{'angle': t_angle, 'distance': t_distance}}})
              outputs = self.envs.step(env_actions)
              observations, _, dones, infos = [list(x) for x in zip(*outputs)]
  
              # calculate reward
              if train_rl:
                  dist = np.zeros(self.envs.num_envs, np.float32)
                  ndtw_score = np.zeros(self.envs.num_envs, np.float32)
                  reward = np.zeros(self.envs.num_envs, np.float32)
                  last_dist_ = last_dist[not_done_index].copy()
                  last_ndtw_ = last_ndtw[not_done_index].copy()
                  not_done_index_bank.append(deepcopy(not_done_index))
                  for i, info in enumerate(infos):
                      dist[i] = info['distance_to_goal']
                      ndtw_score[i] = info['ndtw']
                      if env_actions[i]['action']['action'] == 0: #stop
                          if dist[i] < 3.0:
                              reward[i] = 2.0 + ndtw_score[i] * 2.0
                          else:
                              reward[i] = -2.0
                      else:
                          reward[i] = - (dist[i] - last_dist_[i])
                          ndtw_reward = ndtw_score[i] - last_ndtw_[i]
                          if reward[i] > 0.0:                          
                                  reward[i] = 1.0 + ndtw_reward
                          elif reward[i] < 0.0:
                              reward[i] = -1.0 + ndtw_reward
                          else:
                              reward[i] = -0.5  # get stuck
  
                          if (last_dist_[i] <= 1.0) and (dist[i]-last_dist_[i] > 0.0):
                              reward[i] -= (1.0 - last_dist_[i]) * 2.0
                  rewards.append(reward)
                  last_dist[not_done_index] = dist.copy()
                  last_ndtw[not_done_index] = ndtw_score.copy()
  
              # pause env
              if sum(dones) > 0:
                  for i in reversed(list(range(self.envs.num_envs))):
                      if dones[i]:
                          not_done_index.pop(i)
                          self.envs.pause_at(i)
                          observations.pop(i)
                          # states.pop(i)
                          video_inputs.pop(i)
                          # images_to_video([im for im in history_images[i][k]],'hamt_vis/','history_%s'%(i),4)
                          history_images.pop(i)
              if self.envs.num_envs == 0:
                  break
              
              for i in range(len(observations)):
                  states_i = observations[i].pop('states')
                  # states[i] += states_i
                  new_images_i = {k:[] for k in keys}
                  for position, rotation in states_i[:-1]:
                      new_image = self.envs.call_at(i,'get_pano_rgbs_observations_at', {'source_position':position,'source_rotation':rotation})
                      for k in keys:
                          new_images_i[k].append(new_image[k][None,...])
                  for k in keys:
                      new_images_i[k].append(observations[i][k][None,...])
                      history_images[i][k] = np.vstack((history_images[i][k], np.vstack(new_images_i[k])))
                      if len(history_images[i][k]) < 16:
                          video_inputs[i][k][16-len(history_images[i][k]):] = history_images[i][k]
                      else:
                          video_inputs[i][k] = history_images[i][k][-16:]
                  # print(i,stepk,len(new_images_i[k]))
                  position, rotation = states_i[-1]
                  self.envs.call_at(i,'set_agent_state', {'position':position,'rotation':rotation})
              
              # for i in range(len(observations)):
              #     # for k in video_inputs[i].keys():
              #         k = 'rgb'
              #         images_to_video([im for im in video_inputs[i][k]],'hamt_vis/','%s_%s_%s'%(i,stepk,k),4)
  
              hist_lens = hist_lens[np.array(dones)==False]
              for j in range(len(hist_embeds)):
                  hist_embeds[j] = hist_embeds[j][np.array(dones)==False]
              observations = extract_instruction_tokens(observations,self.config.TASK_CONFIG.TASK.INSTRUCTION_SENSOR_UUID)
              batch = batch_obs(observations, self.device)
              batch = apply_obs_transforms_batch(batch, self.obs_transforms)
              batch['video_rgbs'] = video_inputs
  
          if train_rl:
              rl_loss = 0.
              discount_reward = np.zeros(init_num_envs, np.float32)
              max_length = len(rewards)
              for t in range(max_length-1, -1, -1):
                  index_ = not_done_index_bank[t]
                  discount_reward[index_] = discount_reward[index_] * 0.9 + rewards[t]
                  clip_reward = discount_reward[index_].copy()
                  r_ = torch.from_numpy(clip_reward).cuda()
                  v_ = self.policy.net(mode='critic', critic_states=hidden_states[t])
                  a_ = (r_ - v_).detach()
  
                  t_policy_loss = (-policy_log_probs[t] * a_).sum()
                  t_critic_loss = ((r_ - v_) ** 2).sum() * 0.5 # 1/2 L2 loss
                  t_entropy_loss = -0.01 * entropys[t].sum()
                  rl_loss += (t_policy_loss + t_critic_loss + t_entropy_loss)
  
                  self.logs['policy_loss'].append(t_policy_loss.mean().item())
                  self.logs['critic_loss'].append(t_critic_loss.mean().item())
              rl_loss /= total_actions
              self.loss += rl_loss
              self.logs['RL_loss'].append(rl_loss.item())
          
          if train_ml is not None:
              il_loss = train_ml * il_loss / total_actions
              self.loss += il_loss
              self.logs['IL_loss'].append(il_loss.item())
      
      def _train_interval(self, interval, ml_weight=1., sample_ratio=None, feedback='teacher'):
          self.policy.train()
          if self.world_size > 1:
              self.policy.net.module.rgb_encoder.eval()
              self.policy.net.module.depth_encoder.eval()
          else:
              self.policy.net.rgb_encoder.eval()
              self.policy.net.depth_encoder.eval()
          self.waypoint_predictor.eval()
  
          if self.local_rank < 1:
              pbar = tqdm.trange(interval, leave=False, dynamic_ncols=True)
          else:
              pbar = range(interval)
          self.logs = defaultdict(list)
  
          for idx in pbar:
              self.optimizer.zero_grad()
              self.loss = 0.
  
              if feedback == 'teacher':
                  self._train_iter(train_ml=1., sample_ratio=sample_ratio, train_rl=False)
              elif feedback == 'sample':
                  if ml_weight != 0:
                      self._train_iter(train_ml=ml_weight, sample_ratio=sample_ratio, train_rl=False)
                  self._train_iter(train_ml=None, sample_ratio=None, train_rl=True)
              else:
                  raise NotImplementedError
  
              self.scaler.scale(self.loss).backward() #self.loss.backward()
              self.scaler.unscale_(self.optimizer)
              if self.world_size > 1: 
                  torch.nn.utils.clip_grad_norm_(self.policy.net.module.vln_bert.parameters(), 40.)
              else:
                  torch.nn.utils.clip_grad_norm_(self.policy.net.vln_bert.parameters(), 40.)
              self.scaler.step(self.optimizer) #self.optimizer.step()
              self.scaler.update()
  
              if self.local_rank < 1:
                  pbar.set_postfix({'iter': f'{idx+1}/{interval}'})
              
          return deepcopy(self.logs)
  
      def train(self) -> None:
          self._set_config()
          if self.config.MODEL.task_type == 'rxr':
              self.gt_data = {}
              for role in self.config.TASK_CONFIG.DATASET.ROLES:
                  with gzip.open(
                      self.config.TASK_CONFIG.TASK.NDTW.GT_PATH.format(
                          split=self.split, role=role
                      ), "rt") as f:
                      self.gt_data.update(json.load(f))
  
          observation_space, action_space = self._init_envs()
          start_iter = self._initialize_policy(
              self.config,
              self.config.IL.load_from_ckpt,
              observation_space=observation_space,
              action_space=action_space,
          )
  
          total_iter = self.config.IL.iters
          log_every  = self.config.IL.log_every
          writer     = TensorboardWriter(self.config.TENSORBOARD_DIR if self.local_rank < 1 else None)
  
          self.scaler = GradScaler()
          logger.info('Traning Starts... GOOD LUCK!')
  