plant_agent.py

"""
Privileged driving agent used for data collection.
Drives by accessing the simulator directly.
"""

import os
import torch
import torch.nn.functional as F
from plant import PlanT
from data_agent import DataAgent
from data import CARLA_Data
import cv2
import numpy as np
import jsonpickle
import jsonpickle.ext.numpy as jsonpickle_numpy

import carla
from config import GlobalConfig
import transfuser_utils as t_u

SAVE_PATH = os.environ.get('SAVE_PATH', None)
jsonpickle_numpy.register_handlers()
jsonpickle.set_encoder_options('json', sort_keys=True, indent=4)


def get_entry_point():
  return 'PlanTAgent'


class PlanTAgent(DataAgent):
  """
    Privileged driving agent used for data collection.
    Drives by accessing the simulator directly.
    """

  def setup(self, path_to_conf_file, route_index=None, traffic_manager=None):
    super().setup(path_to_conf_file, route_index, traffic_manager)

    torch.cuda.empty_cache()

    # Load the config saved during training
    with open(os.path.join(path_to_conf_file, 'config.json'), 'rt', encoding='utf-8') as f:
      json_config = f.read()

    loaded_config = jsonpickle.decode(json_config)

    # Generate new config for the case that it has new variables.
    self.config = GlobalConfig()
    # Overwrite all properties that were set in the save config.
    self.config.__dict__.update(loaded_config.__dict__)

    self.config.debug = int(os.environ.get('VISU_PLANT', 0)) == 1
    self.device = torch.device('cuda:0')

    self.data = CARLA_Data(root=[], config=self.config, shared_dict=None)

    self.config.inference_direct_controller = int(os.environ.get('DIRECT', 0))
    self.uncertainty_weight = int(os.environ.get('UNCERTAINTY_WEIGHT', 1))
    self.tuned_aim_distance = int(os.environ.get('TUNED_AIM_DISTANCE', 1))
    print('TUNED_AIM_DISTANCE for wp rep?: ', self.tuned_aim_distance)
    self.compile = int(os.environ.get('COMPILE', 0)) == 1
    print('Uncertainty weighting?: ', self.uncertainty_weight)
    self.config.brake_uncertainty_threshold = float(
        os.environ.get('UNCERTAINTY_THRESHOLD', self.config.brake_uncertainty_threshold))
    if self.uncertainty_weight:
      print('Uncertainty threshold: ', self.config.brake_uncertainty_threshold)

    # Load model files
    self.nets = []
    self.model_count = 0  # Counts how many models are in our ensemble
    for file in os.listdir(path_to_conf_file):
      if file.endswith('.pth') and file.startswith('model'):
        self.model_count += 1
        print(os.path.join(path_to_conf_file, file))
        net = PlanT(self.config)
        if self.config.sync_batch_norm:
          # Model was trained with Sync. Batch Norm.
          # Need to convert it otherwise parameters will load wrong.
          net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)
        state_dict = torch.load(os.path.join(path_to_conf_file, file), map_location=self.device)

        net.load_state_dict(state_dict, strict=False)
        net.cuda()
        net.eval()

        if self.config.compile or self.compile:
          net = torch.compile(net, mode=self.config.compile_mode)

        self.nets.append(net)

    if self.config.debug:
      self.init_map = False

  def sensors(self):
    result = super().sensors()
    if self.config.debug:
      result += [{
          'type': 'sensor.camera.rgb',
          'x': self.config.camera_pos[0],
          'y': self.config.camera_pos[1],
          'z': self.config.camera_pos[2],
          'roll': self.config.camera_rot_0[0],
          'pitch': self.config.camera_rot_0[1],
          'yaw': self.config.camera_rot_0[2],
          'width': self.config.camera_width,
          'height': self.config.camera_height,
          'fov': self.config.camera_fov,
          'id': 'rgb_debug'
      }]
    return result

  @torch.inference_mode()
  def run_step(self, input_data, timestamp, sensors=None):  # pylint: disable=locally-disabled, unused-argument
    if self.config.debug and not self.init_map:
      self.nets[0].init_visualization()
      self.init_map = True

    tick_data = super().run_step(input_data, timestamp, plant=True)

    if self.config.debug:
      camera = input_data['rgb_debug'][1][:, :, :3]
      rgb_debug = cv2.cvtColor(camera, cv2.COLOR_BGR2RGB)
      rgb_debug = np.transpose(rgb_debug, (2, 0, 1))

    target_point = torch.tensor(tick_data['target_point'], dtype=torch.float32).to(self.device).unsqueeze(0)

    # Preprocess route the same way we did during training
    route = tick_data['route']
    if len(route) < self.config.num_route_points:
      num_missing = self.config.num_route_points - len(route)
      route = np.array(route)
      # Fill the empty spots by repeating the last point.
      route = np.vstack((route, np.tile(route[-1], (num_missing, 1))))
    else:
      route = np.array(route[:self.config.num_route_points])

    if self.config.smooth_route:
      route = self.data.smooth_path(route)
    route = torch.tensor(route, dtype=torch.float32)[:self.config.num_route_points].to(self.device).unsqueeze(0)

    light_hazard = torch.tensor(tick_data['light_hazard'], dtype=torch.int32).to(self.device).unsqueeze(0).unsqueeze(0)
    stop_sign_hazard = torch.tensor(tick_data['stop_sign_hazard'],
                                    dtype=torch.int32).to(self.device).unsqueeze(0).unsqueeze(0)
    junction = torch.tensor(tick_data['junction'], dtype=torch.int32).to(self.device).unsqueeze(0).unsqueeze(0)

    bounding_boxes, _ = self.data.parse_bounding_boxes(tick_data['bounding_boxes'])
    bounding_boxes_padded = torch.zeros((self.config.max_num_bbs, 8), dtype=torch.float32).to(self.device)

    if len(bounding_boxes) > 0:
      # Pad bounding boxes to a fixed number
      bounding_boxes = np.stack(bounding_boxes)
      bounding_boxes = torch.tensor(bounding_boxes, dtype=torch.float32).to(self.device)

      if bounding_boxes.shape[0] <= self.config.max_num_bbs:
        bounding_boxes_padded[:bounding_boxes.shape[0], :] = bounding_boxes
      else:
        bounding_boxes_padded[:self.config.max_num_bbs, :] = bounding_boxes[:self.config.max_num_bbs]

    bounding_boxes_padded = bounding_boxes_padded.unsqueeze(0)

    speed = torch.tensor(tick_data['speed'], dtype=torch.float32).to(self.device).unsqueeze(0)

    pred_wps = []
    pred_target_speeds = []
    pred_checkpoints = []
    pred_bbs = []
    for i in range(self.model_count):
      pred_wp, pred_target_speed, pred_checkpoint, pred_bb = self.nets[i].forward(bounding_boxes=bounding_boxes_padded,
                                                                                  route=route,
                                                                                  target_point=target_point,
                                                                                  light_hazard=light_hazard,
                                                                                  stop_hazard=stop_sign_hazard,
                                                                                  junction=junction,
                                                                                  velocity=speed.unsqueeze(1))

      pred_wps.append(pred_wp)
      pred_bbs.append(t_u.plant_quant_to_box(self.config, pred_bb))
      if self.config.use_controller_input_prediction:
        pred_target_speeds.append(F.softmax(pred_target_speed[0], dim=0))
        pred_checkpoints.append(pred_checkpoint[0])

    if self.config.use_wp_gru:
      self.pred_wp = torch.stack(pred_wps, dim=0).mean(dim=0)

    pred_bbs = torch.stack(pred_bbs, dim=0).mean(dim=0)

    uncertainty = None
    if self.config.use_controller_input_prediction:
      pred_target_speed = torch.stack(pred_target_speeds, dim=0).mean(dim=0)

      if self.uncertainty_weight:
        uncertainty = pred_target_speed.detach().cpu().numpy()
        if uncertainty[0] > self.config.brake_uncertainty_threshold:
          pred_target_speed = self.config.target_speeds[0]
        else:
          pred_target_speed = sum(uncertainty * self.config.target_speeds)
      else:
        pred_target_speed_index = torch.argmax(pred_target_speed)
        pred_target_speed = self.config.target_speeds[pred_target_speed_index]

    if self.config.inference_direct_controller and \
        self.config.use_controller_input_prediction:
      pred_checkpoints = torch.stack(pred_checkpoints, dim=0).mean(dim=0).detach().cpu().numpy()
      steer, throttle, brake = self.nets[0].control_pid_direct(pred_checkpoints, pred_target_speed, speed)
    else:
      steer, throttle, brake = self.nets[0].control_pid(self.pred_wp, speed, self.tuned_aim_distance)

    control = carla.VehicleControl()
    control.steer = float(steer)
    control.throttle = float(throttle)
    control.brake = float(brake)

    # Visualize the output of the last model
    if self.config.debug and (not self.save_path is None):
      self.nets[0].visualize_model(save_path=self.save_path,
                                   step=self.step,
                                   rgb=torch.tensor(rgb_debug),
                                   target_point=tick_data['target_point'],
                                   pred_wp=pred_wp,
                                   gt_wp=route,
                                   gt_bbs=bounding_boxes_padded,
                                   pred_speed=uncertainty,
                                   gt_speed=speed,
                                   junction=junction,
                                   light_hazard=light_hazard,
                                   stop_sign_hazard=stop_sign_hazard,
                                   pred_bb=pred_bbs,
                                   pred_checkpoint=pred_checkpoints)

    return control

  def destroy(self, results=None):
    del self.nets
    super().destroy(results)