monoforce_optimization

#! /usr/bin/env python

import os
import torch
import warp as wp
import numpy as np
from matplotlib import pyplot as plt
from tqdm import tqdm
from src.config import DPhysConfig
from src.utils import read_yaml, denormalize_img, ego_to_cam, get_only_in_img_mask
from src.datasets.robingas import compile_data
from src.models.TerrainEncoder import compile_model
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
from datetime import datetime
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from src.models.DiffSim import DiffSim

wp.init()  # init warp!


torch.manual_seed(42)
torch.cuda.manual_seed(42)
np.random.seed(42)


class Learner:
    def __init__(self, batch_size=1, lr=1e-6, weight_decay=1e-7,
                 ls_cfg_path='config/lss_cfg.yaml',
                 lss_weights_path=None,
                 dphys_cfg_path='config/dphys_cfg.yaml',
                 T_horizon=10.0,
                 dt=0.0005,
                 device='cpu',
                 use_renderer=False,
                 use_cuda_graph=False,
                 small_data=False):
        self.device = device
        self.use_cuda_graph = use_cuda_graph
        self.use_renderer = use_renderer
        self.batch_size = batch_size
        self.T_horizon = T_horizon
        self.dt = dt

        self.lss_cfg = read_yaml(ls_cfg_path)
        self.lss_weights = lss_weights_path
        self.terrain_encoder = self.init_terrain_encoder(self.lss_cfg, weights=self.lss_weights)

        self.dphys_cfg = DPhysConfig()
        self.dphys_cfg.from_yaml(dphys_cfg_path)
        self.dphys = self.init_diff_physics()

        self.geom_hm_loss_fn = torch.nn.MSELoss(reduction='none')
        self.optimizer = torch.optim.Adam(self.terrain_encoder.parameters(), lr=lr, weight_decay=weight_decay)

        self.train_data_loader, self.val_data_loader = self.init_dataloaders(batch_size=batch_size, small=small_data)

        log_dir = f'config/tb_runs/{datetime.now().strftime("%Y_%m_%d_%H_%M_%S")}'
        self.tb_writer = SummaryWriter(log_dir=log_dir)
        self.counter = 0

    def init_dataloaders(self, batch_size=1, small=False):
        train_ds, val_ds = compile_data(robot='tradr', T_horizon=self.T_horizon, dt=self.dt, small=small)
        train_data_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
        val_data_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False)
        return train_data_loader, val_data_loader

    def init_terrain_encoder(self, terrain_encoder_cfg, weights=None):
        terrain_encoder = compile_model(terrain_encoder_cfg['grid_conf'], terrain_encoder_cfg['data_aug_conf'])
        if weights is not None and os.path.exists(weights):
            print(f'Loading pretrained LSS weights from {weights}')
            terrain_encoder.load_state_dict(torch.load(weights, map_location=self.device))
        terrain_encoder.to(self.device)
        terrain_encoder.train()
        return terrain_encoder

    def init_diff_physics(self):
        torch_hms, res = self.get_initial_heightmaps(self.batch_size)
        dphys = DiffSim(torch_hms, res, dt=self.dt, use_renderer=self.use_renderer, device=self.device)
        return dphys

    def get_initial_heightmaps(self, batch_size):
        xbound = self.lss_cfg['grid_conf']['xbound']
        ybound = self.lss_cfg['grid_conf']['ybound']
        grid_res = xbound[2]
        shp = (int((xbound[1] - xbound[0]) / grid_res), int((ybound[1] - ybound[0]) / grid_res))
        torch_hms = [torch.zeros(shp, dtype=torch.float32, device=self.device, requires_grad=True)
                     for _ in range(batch_size)]
        res = [grid_res for _ in range(batch_size)]  # heightmap resolutions
        return torch_hms, res

    def geom_hm_loss(self, height_pred, height_gt, weights=None):
        assert height_pred.shape == height_gt.shape, 'Height prediction and ground truth must have the same shape'
        if weights is None:
            weights = torch.ones_like(height_gt)
        assert weights.shape == height_gt.shape, 'Weights and height ground truth must have the same shape'

        # handle imbalanced height distribution (increase weights for higher heights / obstacles)
        h_mean = height_gt[weights.bool()].mean()
        # the higher the difference from mean the higher the weight
        weights_h = 1.0 + torch.abs(height_gt - h_mean)
        # apply height difference weights
        weights = weights * weights_h

        # compute weighted loss
        loss = (self.geom_hm_loss_fn(height_pred * weights, height_gt * weights)).mean()
        return loss

    def traj_loss(self, height, timestamps=None):
        _, loss_traj = self.dphys.simulate_and_backward_torch_tensor(height, use_graph=self.use_cuda_graph)
        loss_traj = wp.to_torch(loss_traj) / self.batch_size
        if timestamps is not None and timestamps.shape[1] > 1:
            loss_traj /= timestamps.shape[1]
        return loss_traj

    def step(self, batch, is_train=True):
        batch = [torch.as_tensor(b, dtype=torch.float32).to(self.device) for b in batch]
        (imgs, rots, trans, intrins, post_rots, post_trans,
         hm_geom,
         timestamps, poses, controls) = batch
        imgs_data = [imgs, rots, trans, intrins, post_rots, post_trans]

        # predict heightmaps
        voxel_feats = self.terrain_encoder.get_voxels(*imgs_data)
        height_pred_geom, height_pred_diff = self.terrain_encoder.bevencode(voxel_feats)
        height_pred_terrain = height_pred_geom - height_pred_diff
        height_pred_geom = height_pred_geom.squeeze(1)
        height_pred_terrain = height_pred_terrain.squeeze(1)
        assert height_pred_geom.shape == hm_geom[:, 0].shape, 'Height prediction and ground truth must have the same shape'

        # dphysics trajectory loss
        loss_traj = self.traj_loss(height_pred_terrain, timestamps)

        # geometry heightmap loss
        loss_geom = self.geom_hm_loss(height_gt=hm_geom[:, 0], height_pred=height_pred_geom, weights=hm_geom[:, 1])

        print(f'Loss traj: {loss_traj.item()}')
        print(f'Loss geom: {loss_geom.item()}')
        # print('HM grad mean:', height_pred.grad.abs().mean().item())

        mode = 'train' if is_train else 'val'
        self.tb_writer.add_scalar(f'{mode}/loss/traj', loss_traj, self.counter)
        self.tb_writer.add_scalar(f'{mode}/loss/geom', loss_geom, self.counter)
        self.tb_writer.add_scalar(f'{mode}/loss/total', loss_traj + loss_geom, self.counter)
        self.tb_writer.flush()

        if is_train:
            height_pred_terrain.backward(height_pred_terrain.grad, retain_graph=True)
            loss_geom.backward(retain_graph=True)

            torch.nn.utils.clip_grad_norm_(self.terrain_encoder.parameters(), max_norm=5.0)
            self.optimizer.step()

            # necessary, since tape.zero() does not reach the torch tensor for some reason
            self.optimizer.zero_grad(set_to_none=False)

        return loss_traj, loss_geom

    def vis_pred(self, batch):
        fig = plt.figure(figsize=(20, 12))
        ax1 = fig.add_subplot(341)
        ax2 = fig.add_subplot(342)
        ax3 = fig.add_subplot(343)
        ax4 = fig.add_subplot(344)
        ax5 = fig.add_subplot(345)
        ax6 = fig.add_subplot(346)
        ax7 = fig.add_subplot(347)
        ax8 = fig.add_subplot(348)

        axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8]
        for ax in axes:
            ax.clear()

        # visualize training predictions
        with torch.no_grad():
            (imgs, rots, trans, intrins, post_rots, post_trans,
             hm_geom,
             timestamps, poses, controls) = batch
            inputs = [imgs, rots, trans, intrins, post_rots, post_trans]
            inputs = [torch.as_tensor(i, dtype=torch.float32).to(self.device) for i in inputs]
            voxel_feats = self.terrain_encoder.get_voxels(*inputs)
            height_pred_geom, height_pred_diff = self.terrain_encoder.bevencode(voxel_feats)
            height_pred_terrain = height_pred_geom - height_pred_diff

            batchi = 0
            height_pred_geom = height_pred_geom[batchi, 0].cpu()
            height_pred_terrain = height_pred_terrain[batchi, 0].cpu()
            height_pred_diff = height_pred_diff[batchi, 0].cpu()
            height_geom = hm_geom[batchi, 0].cpu()

            # get height map points
            z_grid = height_pred_terrain
            x_grid = torch.arange(-self.dphys_cfg.d_max, self.dphys_cfg.d_max, self.dphys_cfg.grid_res)
            y_grid = torch.arange(-self.dphys_cfg.d_max, self.dphys_cfg.d_max, self.dphys_cfg.grid_res)
            x_grid, y_grid = torch.meshgrid(x_grid, y_grid)
            hm_points = torch.stack([x_grid, y_grid, z_grid], dim=-1)
            hm_points = hm_points.view(-1, 3).T

            # plot images with projected height map points
            for imgi in range(imgs.shape[1])[:4]:
                ax = axes[imgi]
                img = imgs[batchi, imgi]
                img = denormalize_img(img[:3])
                cam_pts = ego_to_cam(hm_points, rots[batchi, imgi], trans[batchi, imgi], intrins[batchi, imgi])
                img_H, img_W = self.lss_cfg['data_aug_conf']['H'], self.lss_cfg['data_aug_conf']['W']
                mask_img = get_only_in_img_mask(cam_pts, img_H, img_W)
                plot_pts = post_rots[batchi, imgi].matmul(cam_pts) + post_trans[batchi, imgi].unsqueeze(1)
                ax.imshow(img)
                ax.scatter(plot_pts[0, mask_img], plot_pts[1, mask_img], s=1, c=hm_points[2, mask_img],
                           cmap='jet', vmin=-1.0, vmax=1.0)
                ax.axis('off')

            ax5.set_title('Prediction: Geom')
            ax5.imshow(height_pred_geom.T, origin='lower', cmap='jet', vmin=-1.0, vmax=1.0)

            ax6.set_title('Label: Geom')
            ax6.imshow(height_geom.T, origin='lower', cmap='jet', vmin=-1.0, vmax=1.0)

            ax7.set_title('Prediction: Terrain')
            ax7.imshow(height_pred_terrain.T, origin='lower', cmap='jet', vmin=-1.0, vmax=1.0)

            ax8.set_title('Prediction: HM Diff')
            ax8.imshow(height_pred_diff.T, origin='lower', cmap='jet', vmin=-1.0, vmax=1.0)

            return fig

    def set_poses_and_controls(self, batch):
        (imgs, rots, trans, intrins, post_rots, post_trans,
         hm_geom,
         timestamps, poses, controls) = batch
        flipper_angles = torch.zeros((controls.shape[1], self.batch_size, 4))
        poses[:, :, 2] += 1.0  # add height offset

        self.dphys.set_T(T=int(self.T_horizon / self.dt), T_s=300)
        self.dphys.set_target_poses(timestamps * 1000, poses)
        self.dphys.set_controls(controls, flipper_angles)

        if self.use_renderer:
            self.dphys.render_heightmaps()
            self.dphys.render_traj(poses[0, :, :3].cpu().numpy())

    def test_control(self):
        for batch in tqdm(self.train_data_loader):

            # # show front image
            # img = denormalize_img(batch[0][0, 0, :3])
            # plt.imshow(img)
            # plt.show()

            # show control sequence
            self.set_poses_and_controls(batch)
            self.dphys.init_shoot_states()  # load initial states for the shooter
            self.dphys.save_shoot_init_vels()  # save states for shooter
            self.dphys.simulate_single()  # simulate a single long trajectory for testing
            self.dphys.render_states('current', color=(1.0, 0.0, 0.0))
            self.dphys.render_simulation(pause=False)

            # trajectory loss
            height0, _ = self.get_initial_heightmaps(self.batch_size)
            height0 = torch.stack(height0, dim=0).to(self.device)
            timestamps = batch[7]
            loss_traj = self.traj_loss(height0, timestamps)
            print(f'Loss traj: {loss_traj.item()}')

    def epoch(self, data_loader, is_train=True):
        loss = 0.0
        for batch in tqdm(data_loader):
            self.set_poses_and_controls(batch)
            self.dphys.init_shoot_states()  # load initial states for the shooter

            loss_traj, loss_geom = self.step(batch, is_train=is_train)
            # update loss
            loss += (loss_traj + loss_geom).item()
            # update counter
            self.counter += 1

        fig = self.vis_pred(batch=next(iter(data_loader)))
        self.tb_writer.add_figure(f"{'train' if is_train else 'val'}/prediction", fig, self.counter)
        self.tb_writer.flush()

        if self.use_renderer:
            self.dphys.save_shoot_init_vels()  # save states for shooter
            self.dphys.simulate_single()  # simulate a single long trajectory for testing
            self.dphys.render_states('current', color=(1.0, 0.0, 0.0))
            self.dphys.render_simulation(pause=False)

        loss /= len(data_loader)
        return loss

    def train(self, n_epochs=1):
        for epoch in range(n_epochs):
            print(f'Epoch {epoch} \n-----------')
            train_loss = self.epoch(self.train_data_loader, is_train=True)
            val_loss = self.epoch(self.val_data_loader, is_train=False)

            self.tb_writer.add_scalar('train/loss/epoch', train_loss, epoch)
            self.tb_writer.add_scalar('val/loss/epoch', val_loss, epoch)
            self.tb_writer.flush()


def main():
    device = "cuda"
    use_renderer = False
    small_data = False
    batch_size = 1
    n_epochs = 10
    lr = 1e-6
    robot = 'tradr'
    # lss_weights_path = f'config/weights/lss_robingas_{robot}.pt'
    lss_weights_path = None
    lss_cfg_path = f'config/lss_cfg_{robot}.yaml'
    dphys_cfg_path = 'config/dphys_cfg.yaml'
    T_horizon = 6.0
    dt = 0.0002

    learner = Learner(batch_size=batch_size,
                      lr=lr,
                      ls_cfg_path=lss_cfg_path,
                      lss_weights_path=lss_weights_path,
                      dphys_cfg_path=dphys_cfg_path,
                      T_horizon=T_horizon,
                      dt=dt,
                      device=device,
                      use_renderer=use_renderer,
                      small_data=small_data)
    learner.train(n_epochs=n_epochs)
    # learner.test_control()


if __name__ == '__main__':
    main()