private scenes

Dockerfile

@@ -15,9 +15,9 @@ ENV PATH /app/miniconda/bin:$PATH
 
 SHELL ["conda", "run","--no-capture-output", "-p","/app/env", "/bin/bash", "-c"]
 
-COPY --chown=1000:1000 ./web_server.py /app/web_server.py
-COPY --chown=1000:1000 ./docker/web_server_config/scene-0383-medium-00.yaml  /app/docker/web_server_config/scene-0383-medium-00.yaml
-COPY --chown=1000:1000 ./download_pre_datas.py /app/download_pre_datas.py
+COPY --chown=1000:1000 ./code /app/code
+COPY --chown=1000:1000 ./web_server.py /app/code/web_server.py
+COPY --chown=1000:1000 ./download_pre_datas.py /app/code/download_pre_datas.py
 
 ENV TCNN_CUDA_ARCHITECTURES 75
 
@@ -32,6 +32,6 @@ RUN ./.pixi/envs/default/bin/python3 -m pip install psutil
 
 RUN ./.pixi/envs/default/bin/python3 -m pip install moviepy
 
-RUN ./.pixi/envs/default/bin/python /app/download_pre_datas.py
+RUN ./.pixi/envs/default/bin/python /app/code/download_pre_datas.py
 
-CMD ["./.pixi/envs/default/bin/python", "web_server.py"]
+CMD ["./.pixi/envs/default/bin/python", "/app/code/web_server.py"]

code/gaussian_renderer/__init__.py

@@ -0,0 +1,231 @@
+import torch
+from scene.gaussian_model import GaussianModel
+from scene.ground_model import GroundModel
+from gsplat.rendering import rasterization
+import roma
+from scene.cameras import Camera
+from torch import Tensor
+
+def euler2matrix(yaw):
+    return torch.tensor([
+        [torch.cos(-yaw), 0, torch.sin(-yaw)],
+        [0, 1, 0],
+        [-torch.sin(-yaw), 0, torch.cos(-yaw)]
+    ]).cuda()
+
+def cat_bgfg(bg, fg, only_xyz=False):
+    if only_xyz:
+        if bg.ground_model is None:
+            bg_feats = [bg.get_xyz]
+        else:
+            bg_feats = [bg.get_full_xyz]
+    else:
+        if bg.ground_model is None:
+            bg_feats = [bg.get_xyz, bg.get_opacity, bg.get_scaling, bg.get_rotation, bg.get_features, bg.get_3D_features]
+        else:
+            bg_feats = [bg.get_full_xyz, bg.get_full_opacity, bg.get_full_scaling, bg.get_full_rotation, bg.get_full_features, bg.get_full_3D_features]
+
+    
+    if len(fg) == 0:
+        return bg_feats
+    
+    output = []
+    for fg_feat, bg_feat in zip(fg, bg_feats):
+        if fg_feat is None:
+            output.append(bg_feat)
+        else:
+            if bg_feat.shape[1] != fg_feat.shape[1]:
+                fg_feat = fg_feat[:, :bg_feat.shape[1], :]
+            output.append(torch.cat((bg_feat, fg_feat), dim=0))
+    
+    return output
+
+def concatenate_all(all_fg):
+    output = []
+    for feat in list(zip(*all_fg)):
+        output.append(torch.cat(feat, dim=0))
+    return output
+
+def proj_uv(xyz, cam):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    intr = torch.as_tensor(cam.K[:3, :3]).float().to(device)  # (3, 3)
+    w2c = torch.linalg.inv(cam.c2w)[:3, :]  # (3, 4)
+
+    c_xyz = (w2c[:3, :3] @ xyz.T).T + w2c[:3, 3]
+    i_xyz = (intr @ c_xyz.mT).mT  # (N, 3)
+    uv = i_xyz[:, [1,0]] / i_xyz[:, -1:].clip(1e-3) # (N, 2)
+    return uv
+
+
+def unicycle_b2w(timestamp, model):
+    pred = model(timestamp)
+    if pred is None:
+        return None
+    pred_a, pred_b, pred_v, pitchroll, pred_yaw, pred_h = pred
+    rt = torch.eye(4).float().cuda()
+    rt[:3,:3] = roma.euler_to_rotmat('xzy', [-pitchroll[0]+torch.pi/2, -pitchroll[1]+torch.pi/2, -pred_yaw+torch.pi/2])
+    rt[1, 3], rt[0, 3], rt[2, 3] = pred_h, pred_a, pred_b
+    return rt
+
+
+def render(viewpoint:Camera, prev_viewpoint:Camera, pc:GaussianModel, dynamic_gaussians:dict, 
+            unicycles:dict, bg_color:Tensor, render_optical=False, planning=[]):
+    """
+    Render the scene. 
+    
+    Background tensor (bg_color) must be on GPU!
+    """
+    timestamp = viewpoint.timestamp
+
+    all_fg = [None, None, None, None, None, None]
+    prev_all_fg = [None]
+
+    if unicycles is None or len(unicycles) == 0:
+        track_dict = viewpoint.dynamics
+        if prev_viewpoint is not None:
+            prev_track_dict = prev_viewpoint.dynamics
+    else:
+        track_dict, prev_track_dict = {}, {}
+        for track_id, B2W in viewpoint.dynamics.items():
+            if track_id in unicycles:
+                B2W = unicycle_b2w(timestamp, unicycles[track_id]['model'])
+            track_dict[track_id] = B2W
+            if prev_viewpoint is not None:
+                prev_B2W = unicycle_b2w(prev_viewpoint.timestamp, unicycles[track_id]['model'])
+                prev_track_dict[track_id] = prev_B2W
+    if len(planning) > 0:
+        for plan_id, B2W in planning[0].items():
+            track_dict[plan_id] = B2W
+        if prev_viewpoint is not None:
+            for plan_id, B2W in planning[1].items():
+                prev_track_dict[plan_id] = B2W
+    
+    all_fg, prev_all_fg = [], []
+    for track_id, B2W in track_dict.items():
+        w_dxyz = (B2W[:3, :3] @ dynamic_gaussians[track_id].get_xyz.T).T + B2W[:3, 3]
+
+        drot = roma.quat_wxyz_to_xyzw(dynamic_gaussians[track_id].get_rotation)
+        drot = roma.unitquat_to_rotmat(drot)
+        w_drot = roma.quat_xyzw_to_wxyz(roma.rotmat_to_unitquat(B2W[:3, :3] @ drot))
+        fg = [w_dxyz, 
+            dynamic_gaussians[track_id].get_opacity, 
+            dynamic_gaussians[track_id].get_scaling, 
+            w_drot,
+            # dynamic_gaussians[track_id].get_rotation,
+            dynamic_gaussians[track_id].get_features,
+            dynamic_gaussians[track_id].get_3D_features]
+            
+        all_fg.append(fg)
+
+        if render_optical and prev_viewpoint is not None:
+            if track_id in prev_track_dict:
+                prev_B2W = prev_track_dict[track_id]
+                prev_w_dxyz = torch.mm(prev_B2W[:3, :3], dynamic_gaussians[track_id].get_xyz.T).T + prev_B2W[:3, 3]
+                prev_all_fg.append([prev_w_dxyz])
+            else:
+                prev_all_fg.append([w_dxyz])
+    
+    all_fg = concatenate_all(all_fg)
+    xyz, opacities, scales, rotations, shs, feats3D = cat_bgfg(pc, all_fg)
+
+    if render_optical and prev_viewpoint is not None:
+        prev_all_fg = concatenate_all(prev_all_fg)
+        prev_xyz = cat_bgfg(pc, prev_all_fg, only_xyz=True)[0]
+        uv = proj_uv(xyz, viewpoint)
+        prev_uv = proj_uv(prev_xyz, prev_viewpoint)
+        delta_uv = prev_uv - uv
+        delta_uv = torch.cat([delta_uv, torch.ones_like(delta_uv[:, :1], device=delta_uv.device)], dim=-1)
+    else:
+        delta_uv = torch.zeros_like(xyz)
+
+    if pc.affine:
+        cam_xyz, cam_dir = viewpoint.c2w[:3, 3].cuda(), viewpoint.c2w[:3, 2].cuda()
+        o_enc = pc.pos_enc(cam_xyz[None, :] / 60)
+        d_enc = pc.dir_enc(cam_dir[None, :])
+        appearance = pc.appearance_model(torch.cat([o_enc, d_enc], dim=1)) * 1e-1
+        affine_weight, affine_bias = appearance[:, :9].view(3, 3), appearance[:, -3:]
+        affine_weight = affine_weight + torch.eye(3, device=appearance.device)
+    
+    if render_optical:
+        render_mode = 'RGB+ED+S+F'
+    else:
+        render_mode = 'RGB+ED+S'
+        
+    renders, render_alphas, info = rasterization(
+        means=xyz,
+        quats=rotations,
+        scales=scales,
+        opacities=opacities[:, 0],
+        colors=shs,
+        viewmats=torch.linalg.inv(viewpoint.c2w)[None, ...],  # [C, 4, 4]
+        Ks=viewpoint.K[None, :3, :3],  # [C, 3, 3]
+        width=viewpoint.width,
+        height=viewpoint.height,
+        smts=feats3D[None, ...],
+        flows= delta_uv[None, ...],
+        render_mode=render_mode,
+        sh_degree=pc.active_sh_degree,
+        near_plane=0.01,
+        far_plane=500,
+        packed=False,
+        backgrounds=bg_color[None, :],
+    )
+
+    renders = renders[0]
+    rendered_image = renders[..., :3].permute(2,0,1)
+    depth = renders[..., 3][None, ...]
+    smt = renders[..., 4:(4+feats3D.shape[-1])].permute(2,0,1)
+    
+    if pc.affine:
+        colors = rendered_image.view(3, -1).permute(1, 0) # (H*W, 3)
+        refined_image = (colors @ affine_weight + affine_bias).clip(0, 1).permute(1, 0).view(*rendered_image.shape)
+    else:
+        refined_image = rendered_image
+
+    return {"render": refined_image,
+            "feats": smt,
+            "depth": depth,
+            "opticalflow": renders[..., -2:].permute(2,0,1) if render_optical else None,
+            "alphas": render_alphas,
+            "viewspace_points": info["means2d"],
+            "info": info,
+            }
+
+
+def render_ground(viewpoint:Camera, pc:GroundModel, bg_color:Tensor):
+    xyz, opacities, scales = pc.get_xyz, pc.get_opacity, pc.get_scaling
+    rotations, shs, feats3D = pc.get_rotation, pc.get_features, pc.get_3D_features
+
+    K = viewpoint.K[None, :3, :3]
+    renders, render_alphas, info = rasterization(
+        means=xyz,
+        quats=rotations,
+        scales=scales,
+        opacities=opacities[:, 0],
+        colors=shs,
+        viewmats=torch.linalg.inv(viewpoint.c2w)[None, ...],  # [C, 4, 4]
+        Ks=K,  # [C, 3, 3]
+        width=viewpoint.width,
+        height=viewpoint.height,
+        smts=feats3D[None, ...],
+        render_mode='RGB+ED+S',
+        sh_degree=pc.active_sh_degree,
+        near_plane=0.01,
+        far_plane=500,
+        packed=False,
+        backgrounds=bg_color[None, :],
+    )
+
+    renders = renders[0]
+    rendered_image = renders[..., :3].permute(2,0,1)
+    depth = renders[..., 3][None, ...]
+    smt = renders[..., 4:(4+feats3D.shape[-1])].permute(2,0,1)
+
+    return {"render": rendered_image,
+            "feats": smt,
+            "depth": depth,
+            "opticalflow": None,
+            "alphas": render_alphas,
+            "viewspace_points": info["means2d"],
+            "info": info,
+            }

code/scene/__init__.py

@@ -0,0 +1,111 @@
+import os
+import random
+import json
+from utils.system_utils import searchForMaxIteration
+from scene.dataset_readers import sceneLoadTypeCallbacks
+from scene.gaussian_model import GaussianModel
+from scene.obj_model import ObjModel
+from scene.cameras import cameraList_from_camInfos
+import torch
+import open3d as o3d
+import numpy as np
+import shutil
+
+
+def load_cameras(args, data_type, ignore_dynamic=False):
+    train_cameras = {}
+    test_cameras = {}
+    if os.path.exists(os.path.join(args.source_path, "meta_data.json")):
+        print("Found meta_data.json file, assuming HUGSIM format data set!")
+        scene_info = sceneLoadTypeCallbacks['HUGSIM'](args.source_path, data_type, ignore_dynamic)
+    else:
+        assert False, "Could not recognize scene type! "+args.source_path
+
+    print("Loading Training Cameras")
+    train_cameras = cameraList_from_camInfos(scene_info.train_cameras, args)
+    print("Loading Test Cameras")
+    test_cameras = cameraList_from_camInfos(scene_info.test_cameras, args)
+    return train_cameras, test_cameras, scene_info 
+
+class Scene:
+
+    def __init__(self, args, gaussians:GaussianModel, load_iteration=None, shuffle=True, 
+                 data_type='kitti360', ignore_dynamic=False, planning=None):
+        """b
+        :param path: Path to colmap scene main folder.
+        """
+        self.model_path = args.model_path
+        self.loaded_iter = None
+        self.gaussians = gaussians
+        self.data_type = data_type
+
+        if load_iteration:
+            if load_iteration == -1:
+                self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "ckpts"))
+            else:
+                self.loaded_iter = load_iteration
+            print("Loading trained model at iteration {}".format(self.loaded_iter))
+
+        self.train_cameras, self.test_cameras, scene_info = load_cameras(args, data_type, ignore_dynamic)
+
+        self.dynamic_verts = scene_info.verts
+        self.dynamic_gaussians = {}
+        for track_id in scene_info.verts:
+            self.dynamic_gaussians[track_id] = ObjModel(args.model.sh_degree, feat_mutable=False)
+        if planning is not None:
+            for plan_id in planning.keys():
+                self.dynamic_gaussians[plan_id] = ObjModel(args.model.sh_degree, feat_mutable=False)
+
+        if not self.loaded_iter:
+            shutil.copyfile(scene_info.ply_path, os.path.join(self.model_path, "input.ply"))
+            shutil.copyfile(os.path.join(args.source_path, 'meta_data.json'), os.path.join(self.model_path, 'meta_data.json'))
+            shutil.copyfile(os.path.join(args.source_path, 'ground_param.pkl'), os.path.join(self.model_path, 'ground_param.pkl'))
+
+        if shuffle:
+            random.shuffle(scene_info.train_cameras)
+            random.shuffle(scene_info.test_cameras)
+
+        self.cameras_extent = scene_info.nerf_normalization["radius"]
+
+        if self.loaded_iter:
+            (model_params, first_iter) = torch.load(os.path.join(self.model_path, "ckpts", f"chkpnt{self.loaded_iter}.pth"))
+            gaussians.restore(model_params, None)
+            for iid, dynamic_gaussian in self.dynamic_gaussians.items():
+                if planning is None or iid not in planning:
+                    (model_params, first_iter) = torch.load(os.path.join(self.model_path, "ckpts", f"dynamic_{iid}_chkpnt{self.loaded_iter}.pth"))
+                    dynamic_gaussian.restore(model_params, None)
+                else:
+                    (model_params, first_iter) = torch.load(planning[iid])
+                    model_params = list(model_params)
+                    model_params.append(None)
+                    dynamic_gaussian.restore(model_params, None)
+            # for iid, unicycle_pkg in self.unicycles.items():
+            #     model_params = torch.load(os.path.join(self.model_path, "ckpts", f"unicycle_{iid}_chkpnt{self.loaded_iter}.pth"))
+            #     unicycle_pkg['model'].restore(model_params)
+
+        else:
+            self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent)
+            for track_id in self.dynamic_gaussians.keys():
+                vertices = scene_info.verts[track_id]
+
+                # init from template
+                l, h, w = vertices[:, 0].max() - vertices[:, 0].min(), vertices[:, 1].max() - vertices[:, 1].min(), vertices[:, 2].max() - vertices[:, 2].min()
+                pcd = o3d.io.read_point_cloud(f"utils/vehicle_template/benz_{data_type}.ply")
+                points = np.array(pcd.points) * np.array([l, h, w])
+                pcd.points = o3d.utility.Vector3dVector(points)
+                pcd.colors = o3d.utility.Vector3dVector(np.ones_like(points) * 0.5)
+
+                self.dynamic_gaussians[track_id].create_from_pcd(pcd, self.cameras_extent)
+
+    def save(self, iteration):
+        # self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply"))
+        point_cloud_vis_path = os.path.join(self.model_path, "point_cloud_vis/iteration_{}".format(iteration))
+        self.gaussians.save_vis_ply(os.path.join(point_cloud_vis_path, "point.ply"))
+        for iid, dynamic_gaussian in self.dynamic_gaussians.items():
+            dynamic_gaussian.save_vis_ply(os.path.join(point_cloud_vis_path, f"dynamic_{iid}.ply"))
+
+    def getTrainCameras(self):
+        return self.train_cameras
+
+    def getTestCameras(self):
+        return self.test_cameras

code/scene/cameras.py

@@ -0,0 +1,71 @@
+import torch
+from torch import nn
+
+class Camera(nn.Module):
+    def __init__(self, width, height, image, K, c2w,
+                 image_name, data_device="cuda", 
+                 semantic2d=None, depth=None, mask=None, timestamp=-1, optical_image=None, dynamics={}
+                 ):
+        super(Camera, self).__init__()
+
+        try:
+            self.data_device = torch.device(data_device)
+        except Exception as e:
+            print(e)
+            print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" )
+            self.data_device = torch.device("cuda")
+            
+        self.width = width
+        self.height = height
+        self.image_name = image_name
+        self.timestamp = timestamp
+        self.K = torch.from_numpy(K).float().cuda()
+        self.c2w = torch.from_numpy(c2w).float().cuda()
+        self.dynamics = dynamics
+
+        self.original_image = torch.from_numpy(image).permute(2,0,1).float().clamp(0.0, 1.0).to(self.data_device)
+        if semantic2d is not None:
+            self.semantic2d = semantic2d.to(self.data_device)
+        else:
+            self.semantic2d = None
+        if depth is not None:
+            self.depth = depth.to(self.data_device)
+        else:
+            self.depth = None
+        if mask is not None:
+            self.mask = torch.from_numpy(mask).bool().to(self.data_device)
+        else:
+            self.mask = None
+        self.image_width = self.original_image.shape[2]
+        self.image_height = self.original_image.shape[1]
+        if optical_image is not None:
+            self.optical_gt = torch.from_numpy(optical_image).to(self.data_device)
+        else:
+            self.optical_gt = None
+
+
+def loadCam(args, cam_info):
+
+    if cam_info.semantic2d is not None:
+        semantic2d = torch.from_numpy(cam_info.semantic2d).long()[None, ...]
+    else:
+        semantic2d = None
+
+    optical_image = cam_info.optical_image
+    mask = cam_info.mask
+    depth = cam_info.depth
+
+    gt_image = cam_info.image[..., :3] / 255.
+
+    return Camera(K=cam_info.K, c2w=cam_info.c2w, width=cam_info.width, height=cam_info.height,
+                  image=gt_image, image_name=cam_info.image_name, data_device=args.model.data_device,
+                  semantic2d=semantic2d, depth=depth, mask=mask,
+                  timestamp=cam_info.timestamp, optical_image=optical_image, dynamics=cam_info.dynamics)
+
+def cameraList_from_camInfos(cam_infos, args):
+    camera_list = []
+
+    for c in cam_infos:
+        camera_list.append(loadCam(args, c))
+
+    return camera_list

code/scene/dataset_readers.py

@@ -0,0 +1,212 @@
+import os
+from typing import NamedTuple
+import numpy as np
+import json
+from plyfile import PlyData, PlyElement
+from utils.sh_utils import SH2RGB
+from scene.gaussian_model import BasicPointCloud
+import torch.nn.functional as F
+from imageio.v2 import imread
+import torch
+
+
+class CameraInfo(NamedTuple):
+    K: np.array
+    c2w: np.array
+    image: np.array
+    image_path: str
+    image_name: str
+    width: int
+    height: int
+    semantic2d: np.array
+    optical_image: np.array
+    depth: torch.tensor
+    mask: np.array
+    timestamp: int
+    dynamics: dict
+
+class SceneInfo(NamedTuple):
+    point_cloud: BasicPointCloud
+    train_cameras: list
+    test_cameras: list
+    nerf_normalization: dict
+    ply_path: str
+    verts: dict
+
+def getNerfppNorm(cam_info, data_type):
+    def get_center_and_diag(cam_centers):
+        cam_centers = np.hstack(cam_centers)
+        avg_cam_center = np.mean(cam_centers, axis=1, keepdims=True)
+        center = avg_cam_center
+        dist = np.linalg.norm(cam_centers - center, axis=0, keepdims=True)
+        diagonal = np.max(dist)
+        return center.flatten(), diagonal
+
+    cam_centers = []
+    for cam in cam_info:
+        cam_centers.append(cam.c2w[:3, 3:4]) # cam_centers in world coordinate
+
+    radius = 10
+
+    return {'radius': radius}
+
+def fetchPly(path):
+    plydata = PlyData.read(path)
+    vertices = plydata['vertex']
+    positions = np.vstack([vertices['x'], vertices['y'], vertices['z']]).T
+    if 'red' in vertices:
+        colors = np.vstack([vertices['red'], vertices['green'], vertices['blue']]).T / 255.0
+    else:
+        print('Create random colors')
+        shs = np.ones((positions.shape[0], 3)) * 0.5
+        colors = SH2RGB(shs)
+    normals = np.zeros((positions.shape[0], 3))
+    return BasicPointCloud(points=positions, colors=colors, normals=normals)
+
+def storePly(path, xyz, rgb):
+    # Define the dtype for the structured array
+    dtype = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
+            ('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4'),
+            ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]
+    
+    normals = np.zeros_like(xyz)
+
+    elements = np.empty(xyz.shape[0], dtype=dtype)
+    attributes = np.concatenate((xyz, normals, rgb), axis=1)
+    elements[:] = list(map(tuple, attributes))
+
+    # Create the PlyData object and write to file
+    vertex_element = PlyElement.describe(elements, 'vertex')
+    ply_data = PlyData([vertex_element])
+    ply_data.write(path)
+
+def readHUGSIMCameras(path, data_type, ignore_dynamic):
+    train_cam_infos, test_cam_infos = [], []
+    with open(os.path.join(path, 'meta_data.json')) as json_file:
+        meta_data = json.load(json_file)
+
+        verts = {}
+        if 'verts' in meta_data and not ignore_dynamic:
+            verts_list = meta_data['verts']
+            for k, v in verts_list.items():
+                verts[k] = np.array(v)
+
+        frames = meta_data['frames']
+        for idx, frame in enumerate(frames):
+            c2w = np.array(frame['camtoworld'])
+
+            rgb_path = os.path.join(path, frame['rgb_path'].replace('./', ''))
+
+            rgb_split = rgb_path.split('/')
+            image_name = '_'.join([rgb_split[-2], rgb_split[-1][:-4]])
+            image = imread(rgb_path)
+
+            semantic_2d = None
+            semantic_pth = rgb_path.replace("images", "semantics").replace('.png', '.npy').replace('.jpg', '.npy')
+            if os.path.exists(semantic_pth):
+                semantic_2d = np.load(semantic_pth)
+                semantic_2d[(semantic_2d == 14) | (semantic_2d == 15)] = 13
+
+            optical_path = rgb_path.replace("images", "flow").replace('.png', '_flow.npy').replace('.jpg', '_flow.npy')
+            if os.path.exists(optical_path):
+                optical_image = np.load(optical_path)
+            else:
+                optical_image = None
+
+            depth_path = rgb_path.replace("images", "depth").replace('.png', '.pt').replace('.jpg', '.pt')
+            if os.path.exists(depth_path):
+                depth = torch.load(depth_path, weights_only=True)
+            else:
+                depth = None
+
+            mask = None
+            mask_path = rgb_path.replace("images", "masks").replace('.png', '.npy').replace('.jpg', '.npy')
+            if os.path.exists(mask_path):
+                mask = np.load(mask_path)
+
+            timestamp = frame.get('timestamp', -1)
+
+            intrinsic = np.array(frame['intrinsics'])
+            
+            dynamics = {}
+            if 'dynamics' in frame and not ignore_dynamic:
+                dynamics_list = frame['dynamics']
+                for iid in dynamics_list.keys():
+                    dynamics[iid] = torch.tensor(dynamics_list[iid]).cuda()
+                
+            cam_info = CameraInfo(K=intrinsic, c2w=c2w, image=np.array(image),
+                                image_path=rgb_path, image_name=image_name, height=image.shape[0],
+                                width=image.shape[1], semantic2d=semantic_2d, 
+                                optical_image=optical_image, depth=depth, mask=mask, timestamp=timestamp, dynamics=dynamics)
+            
+            if data_type == 'kitti360':
+                if idx < 20:
+                    train_cam_infos.append(cam_info)
+                elif idx % 20 < 16:
+                    train_cam_infos.append(cam_info)
+                elif idx % 20 >= 16:
+                    test_cam_infos.append(cam_info)
+                else:
+                    continue
+
+            elif data_type == 'kitti':
+                if idx < 10 or idx >= len(frames) - 4:
+                    train_cam_infos.append(cam_info)
+                elif idx % 4 < 2:
+                    train_cam_infos.append(cam_info)
+                elif idx % 4 == 2:
+                    test_cam_infos.append(cam_info)
+                else:
+                    continue
+
+            elif data_type == "nuscenes":
+                if idx % 30 >= 24:
+                    test_cam_infos.append(cam_info)
+                else:
+                    train_cam_infos.append(cam_info)
+
+            elif data_type == "waymo":
+                if idx % 15 >= 12:
+                    test_cam_infos.append(cam_info)
+                else:
+                    train_cam_infos.append(cam_info)
+
+            elif data_type == "pandaset":
+                if idx > 30 and idx % 30 >= 24:
+                    test_cam_infos.append(cam_info)
+                else:
+                    train_cam_infos.append(cam_info)
+            
+            else:
+                raise NotImplementedError
+
+    return train_cam_infos, test_cam_infos, verts
+
+
+def readHUGSIMInfo(path, data_type, ignore_dynamic):
+    train_cam_infos, test_cam_infos, verts = readHUGSIMCameras(path, data_type, ignore_dynamic)
+
+    print(f'Loaded {len(train_cam_infos)} train cameras and {len(test_cam_infos)} test cameras')
+    nerf_normalization = getNerfppNorm(train_cam_infos, data_type)
+
+    ply_path = os.path.join(path, "points3d.ply")
+    if not os.path.exists(ply_path):
+        assert False, "Requires for initialize 3d points as inputs"
+    try:
+        pcd = fetchPly(ply_path)
+    except Exception as e:
+        print('When loading point clound, meet error:', e)
+        exit(0)
+
+    scene_info = SceneInfo(point_cloud=pcd,
+                           train_cameras=train_cam_infos,
+                           test_cameras=test_cam_infos,
+                           nerf_normalization=nerf_normalization,
+                           ply_path=ply_path,
+                           verts=verts)
+    return scene_info
+
+
+sceneLoadTypeCallbacks = {
+    "HUGSIM": readHUGSIMInfo,
+}

code/scene/gaussian_model.py

@@ -0,0 +1,636 @@
+import torch
+import numpy as np
+from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
+from torch import nn
+import os
+from utils.system_utils import mkdir_p
+from plyfile import PlyData, PlyElement
+from utils.sh_utils import RGB2SH, SH2RGB
+from simple_knn._C import distCUDA2
+from utils.graphics_utils import BasicPointCloud
+from utils.general_utils import strip_symmetric, build_scaling_rotation
+import open3d as o3d
+import tinycudann as tcnn
+from math import sqrt
+from scene.ground_model import GroundModel
+from io import BytesIO
+
+                
+class GaussianModel:
+
+    def setup_functions(self):
+        def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+            L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+            actual_covariance = L @ L.transpose(1, 2)
+            symm = strip_symmetric(actual_covariance)
+            return symm
+        
+        self.scaling_activation = torch.exp
+        self.scaling_inverse_activation = torch.log
+
+        self.covariance_activation = build_covariance_from_scaling_rotation
+
+        self.opacity_activation = torch.sigmoid
+        self.inverse_opacity_activation = torch.logit
+
+        self.rotation_activation = torch.nn.functional.normalize
+
+
+    def __init__(self, sh_degree : int, feat_mutable=True, affine=False, ground_args=None):
+        self.active_sh_degree = 0
+        self.max_sh_degree = sh_degree  
+        self._xyz = torch.empty(0)
+        self._features_dc = torch.empty(0)
+        self._features_rest = torch.empty(0)
+        self._feats3D = torch.empty(0)
+        self._scaling = torch.empty(0)
+        self._rotation = torch.empty(0)
+        self._opacity = torch.empty(0)
+        self.max_radii2D = torch.empty(0)
+        self.xyz_gradient_accum = torch.empty(0)
+        self.denom = torch.empty(0)
+        self.optimizer = None
+        self.percent_dense = 0
+        self.spatial_lr_scale = 0
+        self.feat_mutable = feat_mutable
+        self.setup_functions()
+
+        self.pos_enc = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={"otype": "Frequency", "n_frequencies": 2},
+        )
+        self.dir_enc = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={
+                "otype": "SphericalHarmonics",
+                "degree": 3,
+            },
+        )
+
+        self.affine = affine
+        if affine:
+            self.appearance_model = tcnn.Network(
+                n_input_dims=self.pos_enc.n_output_dims + self.dir_enc.n_output_dims,
+                n_output_dims=12,
+                network_config={
+                    "otype": "FullyFusedMLP",
+                    "activation": "ReLU",
+                    "output_activation": "None",
+                    "n_neurons": 32,
+                    "n_hidden_layers": 2,
+                }
+            )
+        else:
+            self.appearance_model = None
+
+        if ground_args:
+            self.ground_model = GroundModel(sh_degree, model_args=ground_args, finetune=True)
+        else:
+            self.ground_model = None
+
+    def capture(self):
+        if self.ground_model is not None:
+            ground_model_params = self.ground_model.capture()
+        else:
+            ground_model_params = None
+        return (
+            self.active_sh_degree,
+            self._xyz,
+            self._features_dc,
+            self._features_rest,
+            self._feats3D,
+            self._scaling,
+            self._rotation,
+            self._opacity,
+            self.spatial_lr_scale,
+            self.appearance_model.state_dict(),
+            ground_model_params,
+        )
+    
+    def restore(self, model_args, training_args):
+        (self.active_sh_degree, 
+        self._xyz, 
+        self._features_dc, 
+        self._features_rest,
+        self._feats3D,
+        self._scaling, 
+        self._rotation, 
+        self._opacity,
+        self.spatial_lr_scale,
+        appearance_state_dict,
+        ground_model_params,
+        ) = model_args
+        self.appearance_model.load_state_dict(appearance_state_dict, strict=False)
+        if training_args is not None:
+            self.training_setup(training_args)
+        if ground_model_params is not None:
+            self.ground_model = GroundModel(self.max_sh_degree, model_args=ground_model_params)
+        
+    @property
+    def get_scaling(self):
+        return self.scaling_activation(self._scaling)
+        
+    @property
+    def get_full_scaling(self):
+        assert self.ground_model is not None
+        return torch.cat([self.scaling_activation(self._scaling), self.ground_model.get_scaling])
+        
+    @property
+    def get_rotation(self):
+        return self.rotation_activation(self._rotation)
+    
+    @property
+    def get_full_rotation(self):
+        assert self.ground_model is not None
+        return torch.cat([self.rotation_activation(self._rotation), self.ground_model.get_rotation])
+    
+    @property
+    def get_xyz(self):
+        return self._xyz
+        
+    @property
+    def get_full_xyz(self):
+        assert self.ground_model is not None
+        return torch.cat([self._xyz, self.ground_model.get_xyz])
+
+    @property
+    def get_features(self):
+        features_dc = self._features_dc
+        features_rest = self._features_rest
+        return torch.cat((features_dc, features_rest), dim=1)
+    
+    @property
+    def get_full_features(self):
+        assert self.ground_model is not None
+        sh = torch.cat((self._features_dc, self._features_rest), dim=1)
+        return torch.cat([sh, self.ground_model.get_features])
+    
+    @property
+    def get_3D_features(self):
+        return torch.softmax(self._feats3D, dim=-1)
+        
+    @property
+    def get_full_3D_features(self):
+        assert self.ground_model is not None
+        return torch.cat([torch.softmax(self._feats3D, dim=-1), self.ground_model.get_3D_features])
+
+    @property
+    def get_opacity(self):
+        return self.opacity_activation(self._opacity)
+        
+    @property
+    def get_full_opacity(self):
+        assert self.ground_model is not None
+        return torch.cat([self.opacity_activation(self._opacity), self.ground_model.get_opacity])
+    
+    # def get_covariance(self, scaling_modifier = 1):
+    #     return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
+
+    def oneupSHdegree(self):
+        if self.active_sh_degree < self.max_sh_degree:
+            self.active_sh_degree += 1
+
+    def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float):
+        # self.spatial_lr_scale = 1
+        self.spatial_lr_scale = spatial_lr_scale
+        fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
+        fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
+        features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
+        features[:, :3, 0 ] = fused_color
+        features[:, 3:, 1:] = 0.0
+
+        if self.feat_mutable:
+            feats3D = torch.rand(fused_color.shape[0], 20).float().cuda()
+            self._feats3D = nn.Parameter(feats3D.requires_grad_(True))
+        else:
+            feats3D = torch.zeros(fused_color.shape[0], 20).float().cuda()
+            feats3D[:, 13] = 1
+            self._feats3D = feats3D
+
+        print("Number of points at initialization : ", fused_point_cloud.shape[0])
+
+        dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
+        scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
+        rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
+        rots[:, 0] = 1
+
+        opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
+
+        self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))
+        self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
+        self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
+        self._scaling = nn.Parameter(scales.requires_grad_(True))
+        self._rotation = nn.Parameter(rots.requires_grad_(True))
+        self._opacity = nn.Parameter(opacities.requires_grad_(True))
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def training_setup(self, training_args):
+        self.percent_dense = training_args.percent_dense
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+
+        # self.spatial_lr_scale /= 3
+
+        l = [
+            {'params': [self._xyz], 'lr': training_args.position_lr_init*self.spatial_lr_scale, "name": "xyz"},
+            {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
+            {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
+            {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
+            {'params': [self._scaling], 'lr': training_args.scaling_lr*self.spatial_lr_scale, "name": "scaling"},
+            {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"},
+        ]
+
+        if self.affine:
+            l.append({'params': [*self.appearance_model.parameters()], 'lr': 1e-3, "name": "appearance_model"})
+        
+        if self.feat_mutable:
+            l.append({'params': [self._feats3D], 'lr': 1e-2, "name": "feats3D"})
+
+        if self.ground_model is not None:
+            self.ground_optimizer = self.ground_model.optimizer
+        else:
+            self.ground_optimizer = None
+
+        self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)
+        self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,
+                                                    lr_final=training_args.position_lr_final*self.spatial_lr_scale,
+                                                    lr_delay_mult=training_args.position_lr_delay_mult,
+                                                    max_steps=training_args.position_lr_max_steps)
+
+    def update_learning_rate(self, iteration):
+        ''' Learning rate scheduling per step '''
+        for param_group in self.optimizer.param_groups:
+            if param_group["name"] == "xyz":
+                lr = self.xyz_scheduler_args(iteration)
+                param_group['lr'] = lr
+                return lr
+
+    def construct_list_of_attributes(self):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        # All channels except the 3 DC
+        for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
+            l.append('f_dc_{}'.format(i))
+        for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):
+            l.append('f_rest_{}'.format(i))
+        for i in range(self._feats3D.shape[1]):
+            l.append('semantic_{}'.format(i))
+        l.append('opacity')
+        for i in range(self._scaling.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(self._rotation.shape[1]):
+            l.append('rot_{}'.format(i))
+        return l
+
+    def save_ply(self, path=None):
+        mkdir_p(os.path.dirname(path))
+
+        if self.ground_model is not None:
+            xyz = self.get_full_xyz.detach().cpu().numpy()
+            normals = np.zeros_like(xyz)
+            f_dc = torch.cat([self._features_dc, self.ground_model._features_dc]).detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+            f_rest = torch.cat([self._features_rest, self.ground_model._features_rest]).detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+            feats3D = torch.cat([self._feats3D, self.ground_model._feats3D]).detach().cpu().numpy()
+            opacities = torch.cat([self._opacity, self.ground_model._opacity]).detach().cpu().numpy()
+            scale = self.scaling_inverse_activation(self.get_full_scaling).detach().cpu().numpy()
+            rotation = torch.cat([self._rotation, self.ground_model._rotation]).detach().cpu().numpy()
+        else:
+            xyz = self.get_xyz.detach().cpu().numpy()
+            normals = np.zeros_like(xyz)
+            f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+            f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+            feats3D = self._feats3D.detach().cpu().numpy()
+            opacities = self._opacity.detach().cpu().numpy()
+            scale = self.scaling_inverse_activation(self.get_scaling).detach().cpu().numpy()
+            rotation = self._rotation.detach().cpu().numpy()
+
+        dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyz, normals, f_dc, f_rest, feats3D, opacities, scale, rotation), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+        plydata = PlyData([el])
+        if path is not None:
+            plydata.write(path)
+        return plydata
+
+    def save_splat(self, ply_path, splat_path):
+        plydata = self.save_ply(ply_path)
+        vert = plydata["vertex"]
+        sorted_indices = np.argsort(
+            -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
+            / (1 + np.exp(-vert["opacity"]))
+        )
+        buffer = BytesIO()
+        for idx in sorted_indices:
+            v = plydata["vertex"][idx]
+            position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32)
+            scales = np.exp(
+                np.array(
+                    [v["scale_0"], v["scale_1"], v["scale_2"]],
+                    dtype=np.float32,
+                )
+            )
+            rot = np.array(
+                [v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]],
+                dtype=np.float32,
+            )
+            SH_C0 = 0.28209479177387814
+            color = np.array(
+                [
+                    0.5 + SH_C0 * v["f_dc_0"],
+                    0.5 + SH_C0 * v["f_dc_1"],
+                    0.5 + SH_C0 * v["f_dc_2"],
+                    1 / (1 + np.exp(-v["opacity"])),
+                ]
+            )
+            buffer.write(position.tobytes())
+            buffer.write(scales.tobytes())
+            buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
+            buffer.write(
+                ((rot / np.linalg.norm(rot)) * 128 + 128)
+                .clip(0, 255)
+                .astype(np.uint8)
+                .tobytes()
+            )
+        with open(splat_path, "wb") as f:
+            f.write(buffer.getvalue())
+
+    def save_semantic_pcd(self, path):
+        color_dict = {
+            0: np.array([128, 64, 128]),  # Road
+            1: np.array([244, 35, 232]),  # Sidewalk
+            2: np.array([70, 70, 70]),  # Building
+            3: np.array([102, 102, 156]),  # Wall
+            4: np.array([190, 153, 153]),  # Fence
+            5: np.array([153, 153, 153]),  # Pole
+            6: np.array([250, 170, 30]),  # Traffic Light
+            7: np.array([220, 220, 0]),  # Traffic Sign
+            8: np.array([107, 142, 35]),  # Vegetation
+            9: np.array([152, 251, 152]),  # Terrain
+            10: np.array([0, 0, 0]),  # Black (trainId 10)
+            11: np.array([70, 130, 180]),  # Sky
+            12: np.array([220, 20, 60]),  # Person
+            13: np.array([255, 0, 0]),  # Rider
+            14: np.array([0, 0, 142]),  # Car
+            15: np.array([0, 0, 70]),  # Truck
+            16: np.array([0, 60, 100]),  # Bus
+            17: np.array([0, 80, 100]),  # Train
+            18: np.array([0, 0, 230]),  # Motorcycle
+            19: np.array([119, 11, 32])  # Bicycle
+        }
+        semantic_idx = torch.argmax(self.get_full_3D_features, dim=-1, keepdim=True)
+        opacities = self.get_full_opacity[:, 0]
+        mask = ((semantic_idx != 10)[:, 0]) & ((semantic_idx != 8)[:, 0]) & (opacities > 0.2)
+
+        semantic_idx = semantic_idx[mask]
+        semantic_rgb = torch.zeros_like(semantic_idx).repeat(1, 3)
+        for idx in range(20):
+            rgb = torch.from_numpy(color_dict[idx]).to(semantic_rgb.device)[None, :]
+            semantic_rgb[(semantic_idx == idx)[:, 0], :] = rgb
+        semantic_rgb = semantic_rgb.float() / 255.0
+        pcd_xyz = self.get_full_xyz[mask]
+        smt_pcd = o3d.geometry.PointCloud()
+        smt_pcd.points = o3d.utility.Vector3dVector(pcd_xyz.detach().cpu().numpy())
+        smt_pcd.colors = o3d.utility.Vector3dVector(semantic_rgb.detach().cpu().numpy())
+        o3d.io.write_point_cloud(path, smt_pcd)
+
+    def save_vis_ply(self, path):
+        mkdir_p(os.path.dirname(path))
+        xyz = self.get_xyz.detach().cpu().numpy()
+        if self.ground_model:
+            xyz = np.concatenate([xyz, self.ground_model.get_xyz.detach().cpu().numpy()])
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(xyz)
+        colors = SH2RGB(self._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1)
+        if self.ground_model:
+            ground_colors = SH2RGB(self.ground_model._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1)
+            colors = np.concatenate([colors, ground_colors])
+        pcd.colors = o3d.utility.Vector3dVector(colors)
+        o3d.io.write_point_cloud(path, pcd)
+
+    def reset_opacity(self):
+        opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
+        optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
+        self._opacity = optimizable_tensors["opacity"]
+
+    def load_ply(self, path):
+        plydata = PlyData.read(path)
+
+        xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                        np.asarray(plydata.elements[0]["y"]),
+                        np.asarray(plydata.elements[0]["z"])),  axis=1)
+        opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+        features_dc = np.zeros((xyz.shape[0], 3, 1))
+        features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+        features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
+        features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+        extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
+        assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3
+        features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
+        for idx, attr_name in enumerate(extra_f_names):
+            features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
+        # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
+        features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
+
+        scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+        scales = np.zeros((xyz.shape[0], len(scale_names)))
+        for idx, attr_name in enumerate(scale_names):
+            scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
+        rots = np.zeros((xyz.shape[0], len(rot_names)))
+        for idx, attr_name in enumerate(rot_names):
+            rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
+        self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
+        self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
+
+        self.active_sh_degree = self.max_sh_degree
+
+    def replace_tensor_to_optimizer(self, tensor, name):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group["name"] == name:
+                stored_state = self.optimizer.state.get(group['params'][0], None)
+                if stored_state is not None:
+                    stored_state["exp_avg"] = torch.zeros_like(tensor)
+                    stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
+                    del self.optimizer.state[group['params'][0]]
+                    group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                    self.optimizer.state[group['params'][0]] = stored_state
+                    optimizable_tensors[group["name"]] = group["params"][0]
+                else:
+                    group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                    optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def _prune_optimizer(self, mask):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group['name'] == 'appearance_model':
+                continue
+            stored_state = self.optimizer.state.get(group['params'][0], None)
+            if stored_state is not None:
+                stored_state["exp_avg"] = stored_state["exp_avg"][mask]
+                stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def prune_points(self, mask):
+        valid_points_mask = ~mask
+        optimizable_tensors = self._prune_optimizer(valid_points_mask)
+
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        if self.feat_mutable:
+            self._feats3D = optimizable_tensors["feats3D"]
+        else:
+            self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1))
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+
+        self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
+
+        self.denom = self.denom[valid_points_mask]
+        self.max_radii2D = self.max_radii2D[valid_points_mask]
+
+    def cat_tensors_to_optimizer(self, tensors_dict):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group['name'] not in tensors_dict:
+                continue
+            assert len(group["params"]) == 1
+            extension_tensor = tensors_dict[group["name"]]
+            stored_state = self.optimizer.state.get(group["params"][0], None)
+            if stored_state is not None:
+
+                stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
+                stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
+
+                del self.optimizer.state[group["params"][0]]
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                self.optimizer.state[group["params"][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+
+        return optimizable_tensors
+
+    def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation):
+        d = {"xyz": new_xyz,
+        "f_dc": new_features_dc,
+        "f_rest": new_features_rest,
+        "feats3D": new_feats3D,
+        "opacity": new_opacities,
+        "scaling" : new_scaling,
+        "rotation" : new_rotation}
+
+        optimizable_tensors = self.cat_tensors_to_optimizer(d)
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        if self.feat_mutable:
+            self._feats3D = optimizable_tensors["feats3D"]
+        else:
+            self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1))
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
+        n_init_points = self.get_xyz.shape[0]
+        # Extract points that satisfy the gradient condition
+        padded_grad = torch.zeros((n_init_points), device="cuda")
+        padded_grad[:grads.shape[0]] = grads.squeeze()
+        selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
+
+        stds = self.get_scaling[selected_pts_mask].repeat(N,1)
+        means =torch.zeros((stds.size(0), 3),device="cuda")
+        samples = torch.normal(mean=means, std=stds)
+        rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
+        new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
+        new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))
+        new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
+        new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
+        new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
+        new_feats3D = self._feats3D[selected_pts_mask].repeat(N,1)
+        new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacity, new_scaling, new_rotation)
+
+        prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
+        self.prune_points(prune_filter)
+
+    def densify_and_clone(self, grads, grad_threshold, scene_extent):
+        # Extract points that satisfy the gradient condition
+        selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
+        
+        new_xyz = self._xyz[selected_pts_mask]
+        new_features_dc = self._features_dc[selected_pts_mask]
+        new_features_rest = self._features_rest[selected_pts_mask]
+        new_feats3D = self._feats3D[selected_pts_mask]
+        new_opacities = self._opacity[selected_pts_mask]
+        new_scaling = self._scaling[selected_pts_mask]
+        new_rotation = self._rotation[selected_pts_mask]
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation)
+
+    def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size, cam_pos=None):
+        grads = self.xyz_gradient_accum / self.denom
+        grads[grads.isnan()] = 0.0
+
+        self.densify_and_clone(grads, max_grad, extent)
+        self.densify_and_split(grads, max_grad, extent)
+
+        prune_mask = (self.get_opacity < min_opacity).squeeze()
+        if max_screen_size:
+            big_points_vs = self.max_radii2D > max_screen_size
+            if cam_pos is not None:
+                # points_cam_dist = torch.abs(self.get_xyz[:, None, :] - cam_pos[None, ...])
+                # points_cam_nearest_idx = torch.argmin(torch.norm(points_cam_dist, dim=-1), dim=1)
+                # points_cam_dist = points_cam_dist[torch.arange(points_cam_dist.shape[0]), points_cam_nearest_idx, :]
+                # near_mask1 = (points_cam_dist[:, 1] < 5) & (points_cam_dist[:, 0] < 10) & (points_cam_dist[:, 2] < 10)
+                # big_points_ws1 = near_mask1 & (self.get_scaling.max(dim=1).values > 1.0)
+                # near_mask2 = (points_cam_dist[:, 1] < 10) & (points_cam_dist[:, 0] < 20) & (points_cam_dist[:, 2] < 20)
+                # big_points_ws2 = near_mask2 & (self.get_scaling.max(dim=1).values > 3.0)
+                # big_points_ws = (self.get_scaling.max(dim=1).values > 10.0) | big_points_ws1 | big_points_ws2 
+                big_points_ws = self.get_scaling.max(dim=1).values > 10
+                prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+            else:
+                big_points_ws = self.get_scaling.max(dim=1).values > 5
+                prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+        self.prune_points(prune_mask)
+
+        torch.cuda.empty_cache()
+    
+    def add_densification_stats_grad(self, tensor_grad, update_filter):
+        self.xyz_gradient_accum[update_filter] += torch.norm(tensor_grad[update_filter,:2], dim=-1, keepdim=True)
+        self.denom[update_filter] += 1
+        

code/scene/ground_model.py

@@ -0,0 +1,360 @@
+import torch
+import numpy as np
+from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
+from torch import nn
+import os
+from utils.system_utils import mkdir_p
+from plyfile import PlyData, PlyElement
+from utils.sh_utils import RGB2SH, SH2RGB
+from simple_knn._C import distCUDA2
+from utils.graphics_utils import BasicPointCloud
+from utils.general_utils import strip_symmetric, build_scaling_rotation
+import open3d as o3d
+import math
+from utils.graphics_utils import BasicPointCloud
+from utils.sh_utils import RGB2SH
+                
+class GroundModel:
+
+    def setup_functions(self):
+        def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+            L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+            actual_covariance = L @ L.transpose(1, 2)
+            symm = strip_symmetric(actual_covariance)
+            return symm
+        
+        self.scaling_activation = torch.exp
+        self.scaling_inverse_activation = torch.log
+
+        self.covariance_activation = build_covariance_from_scaling_rotation
+
+        self.opacity_activation = torch.sigmoid
+        self.inverse_opacity_activation = torch.logit
+
+        self.rotation_activation = torch.nn.functional.normalize
+
+
+    def __init__(self, sh_degree: int, ground_pcd: BasicPointCloud=None, model_args=None, finetune=False):
+        assert not ((ground_pcd is None) and (model_args is None)), "Need at least one way of initialization"
+        self.active_sh_degree = 0
+        self.max_sh_degree = sh_degree 
+
+        self.scale = 0.1
+        
+        if ground_pcd is not None:
+            self._xyz = nn.Parameter(torch.from_numpy(ground_pcd.points).float().cuda())
+            fused_color = RGB2SH(torch.tensor(np.asarray(ground_pcd.colors)).float().cuda())
+            features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
+            features[:, :3, 0 ] = fused_color
+            features[:, 3:, 1:] = 0.0
+            self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
+            self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
+            
+            self._feats3D = torch.zeros((self._xyz.shape[0], 20)).cuda()
+            self._feats3D[:, 1] = 1
+            self._feats3D = nn.Parameter(self._feats3D)
+            self._rotation = torch.zeros((self._xyz.shape[0], 4)).cuda()
+            self._rotation[:, 0] = 1
+            self._opacity = inverse_sigmoid(torch.ones((self._xyz.shape[0], 1)).cuda() * 0.99)
+            self._scaling = nn.Parameter(torch.ones((self._xyz.shape[0], 2)).float().cuda() * math.log(self.scale))
+
+            self.max_radii2D = torch.zeros((self._xyz.shape[0]), device="cuda")
+            self.percent_dense = 0.01
+            self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+            self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        else:
+            self.restore(model_args)
+
+        if finetune:
+            self.param_groups = [
+                {'params': [self._features_dc], 'lr': 2.5e-3, "name": "f_dc"},
+                {'params': [self._features_rest], 'lr': 2.5e-3 / 20.0, "name": "f_rest"},
+                {'params': [self._feats3D], 'lr': 1e-3, "name": "feats3D"},
+            ]
+        else:
+            self.param_groups = [
+                {'params': [self._xyz], 'lr': 1.6e-4, "name": "xyz"},
+                {'params': [self._features_dc], 'lr': 2.5e-3, "name": "f_dc"},
+                {'params': [self._features_rest], 'lr': 2.5e-3 / 20.0, "name": "f_rest"},
+                {'params': [self._feats3D], 'lr': 1e-2, "name": "feats3D"},
+                {'params': [self._opacity], 'lr': 0.05, "name": "opacity"},
+                {'params': [self._scaling], 'lr': 1e-3, "name": "scaling"},
+            ]
+        self.optimizer = torch.optim.Adam(self.param_groups, lr=0.0, eps=1e-15)
+        self.setup_functions()
+
+    def capture(self):
+        return (
+            self.active_sh_degree,
+            self._xyz,
+            # self._y,
+            # self._z,
+            self._features_dc,
+            self._features_rest,
+            self._feats3D,
+            self._scaling,
+            self._rotation,
+            self._opacity,
+        )
+    
+    def restore(self, model_args):
+        (self.active_sh_degree, 
+        self._xyz,
+        # self._y,
+        # self._z,
+        self._features_dc, 
+        self._features_rest,
+        self._feats3D,
+        self._scaling,
+        self._rotation, 
+        self._opacity) = model_args
+
+    @property
+    def get_scaling(self):
+        scale_y = torch.ones_like(self._xyz[:, 0]) * math.log(0.001)
+        scaling = torch.stack((self._scaling[:, 0], scale_y, self._scaling[:, 1]), dim=1).cuda()
+        # scaling = torch.stack((self._scaling, scale_y, self._scaling), dim=1).cuda()
+        return self.scaling_activation(scaling)
+    
+    @property
+    def get_rotation(self):
+        return self.rotation_activation(self._rotation)
+
+    @property
+    def get_xyz(self):
+        return self._xyz
+        
+    @property
+    def get_features(self):
+        features_dc = self._features_dc
+        features_rest = self._features_rest
+        return torch.cat((features_dc, features_rest), dim=1)
+    
+    @property
+    def get_3D_features(self):
+        return torch.softmax(self._feats3D, dim=-1)
+    
+    @property
+    def get_opacity(self):
+        return self.opacity_activation(self._opacity)
+    
+    def get_covariance(self, scaling_modifier = 1):
+        return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
+
+    def oneupSHdegree(self):
+        if self.active_sh_degree < self.max_sh_degree:
+            self.active_sh_degree += 1
+
+    def construct_list_of_attributes(self):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        # All channels except the 3 DC
+        for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
+            l.append('f_dc_{}'.format(i))
+        for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):
+            l.append('f_rest_{}'.format(i))
+        for i in range(self._feats3D.shape[1]):
+            l.append('semantic_{}'.format(i))
+        l.append('opacity')
+        for i in range(self._scaling.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(self._rotation.shape[1]):
+            l.append('rot_{}'.format(i))
+        return l
+
+    def save_ply(self, path):
+        mkdir_p(os.path.dirname(path))
+
+        xyz = self.get_xyz.detach().cpu().numpy()
+        normals = np.zeros_like(xyz)
+        f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        feats3D = self._feats3D.detach().cpu().numpy()
+        opacities = self._opacity.detach().cpu().numpy()
+        scale = self._scaling.detach().cpu().numpy()
+        rotation = self._rotation.detach().cpu().numpy()
+
+        dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyz, normals, f_dc, f_rest, feats3D, opacities, scale, rotation), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+        PlyData([el]).write(path)
+
+    def save_vis_ply(self, path):
+        mkdir_p(os.path.dirname(path))
+        xyz = self.get_xyz.detach().cpu().numpy()
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(xyz)
+        colors = SH2RGB(self._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1)
+        pcd.colors = o3d.utility.Vector3dVector(colors)
+        o3d.io.write_point_cloud(path, pcd)
+
+    def reset_opacity(self):
+        opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
+        optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
+        self._opacity = optimizable_tensors["opacity"]
+
+    def replace_tensor_to_optimizer(self, tensor, name):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group["name"] == name:
+                stored_state = self.optimizer.state.get(group['params'][0], None)
+                if stored_state is not None:
+                    stored_state["exp_avg"] = torch.zeros_like(tensor)
+                    stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
+                    del self.optimizer.state[group['params'][0]]
+                    group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                    self.optimizer.state[group['params'][0]] = stored_state
+                    optimizable_tensors[group["name"]] = group["params"][0]
+                else:
+                    group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                    optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def _prune_optimizer(self, mask):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group['name'] == 'appearance_model':
+                continue
+            stored_state = self.optimizer.state.get(group['params'][0], None)
+            if stored_state is not None:
+                stored_state["exp_avg"] = stored_state["exp_avg"][mask]
+                stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def prune_points(self, mask):
+        valid_points_mask = ~mask
+        optimizable_tensors = self._prune_optimizer(valid_points_mask)
+
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._feats3D = optimizable_tensors["feats3D"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = self._rotation[0, :].repeat((self._xyz.shape[0], 1))
+
+        self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
+
+        self.denom = self.denom[valid_points_mask]
+        self.max_radii2D = self.max_radii2D[valid_points_mask]
+
+    def cat_tensors_to_optimizer(self, tensors_dict):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group['name'] not in tensors_dict:
+                continue
+            assert len(group["params"]) == 1
+            extension_tensor = tensors_dict[group["name"]]
+            stored_state = self.optimizer.state.get(group["params"][0], None)
+            if stored_state is not None:
+                stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
+                stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
+
+                del self.optimizer.state[group["params"][0]]
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                self.optimizer.state[group["params"][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+
+        return optimizable_tensors
+
+    def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation):
+        d = {"xyz": new_xyz,
+        "f_dc": new_features_dc,
+        "f_rest": new_features_rest,
+        "feats3D": new_feats3D,
+        "opacity": new_opacities,
+        "scaling" : new_scaling}
+
+        optimizable_tensors = self.cat_tensors_to_optimizer(d)
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._feats3D = optimizable_tensors["feats3D"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = self._rotation[0, :].repeat((self._xyz.shape[0], 1))
+
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
+        n_init_points = self.get_xyz.shape[0]
+        # Extract points that satisfy the gradient condition
+        padded_grad = torch.zeros((n_init_points), device="cuda")
+        padded_grad[:grads.shape[0]] = grads.squeeze()
+        selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
+
+        stds = self.get_scaling[selected_pts_mask].repeat(N,1)
+        means =torch.zeros((stds.size(0), 3),device="cuda")
+        samples = torch.normal(mean=means, std=stds)
+        rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
+        new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
+        new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))[:, [0,2]]
+        new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
+        new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
+        new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
+        new_feats3D = self._feats3D[selected_pts_mask].repeat(N,1)
+        new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacity, new_scaling, new_rotation)
+
+        prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
+        self.prune_points(prune_filter)
+
+    def densify_and_clone(self, grads, grad_threshold, scene_extent):
+        # Extract points that satisfy the gradient condition
+        selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
+        
+        new_xyz = self._xyz[selected_pts_mask]
+        new_features_dc = self._features_dc[selected_pts_mask]
+        new_features_rest = self._features_rest[selected_pts_mask]
+        new_feats3D = self._feats3D[selected_pts_mask]
+        new_opacities = self._opacity[selected_pts_mask]
+        new_scaling = self._scaling[selected_pts_mask]
+        new_rotation = self._rotation[selected_pts_mask]
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation)
+
+    def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):
+        grads = self.xyz_gradient_accum / self.denom
+        grads[grads.isnan()] = 0.0
+
+        self.densify_and_clone(grads, max_grad, extent)
+        self.densify_and_split(grads, max_grad, extent)
+
+        prune_mask = (self.get_opacity < min_opacity).squeeze()
+        if max_screen_size:
+            big_points_vs = self.max_radii2D > max_screen_size
+            big_points_ws = self.get_scaling.max(dim=1).values > 1.0
+            prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+        self.prune_points(prune_mask)
+
+        torch.cuda.empty_cache()
+
+    def add_densification_stats(self, viewspace_point_tensor, update_filter):
+        self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True)
+        self.denom[update_filter] += 1
+    
+    def add_densification_stats_grad(self, tensor_grad, update_filter):
+        self.xyz_gradient_accum[update_filter] += torch.norm(tensor_grad[update_filter,:2], dim=-1, keepdim=True)
+        self.denom[update_filter] += 1

code/scene/obj_model.py

@@ -0,0 +1,567 @@
+import torch
+import numpy as np
+from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
+from torch import nn
+import os
+from utils.system_utils import mkdir_p
+from plyfile import PlyData, PlyElement
+from utils.sh_utils import RGB2SH, SH2RGB
+from simple_knn._C import distCUDA2
+from utils.graphics_utils import BasicPointCloud
+from utils.general_utils import strip_symmetric, build_scaling_rotation
+import open3d as o3d
+import tinycudann as tcnn
+from io import BytesIO
+
+                
+class ObjModel:
+
+    def setup_functions(self):
+        def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+            L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+            actual_covariance = L @ L.transpose(1, 2)
+            symm = strip_symmetric(actual_covariance)
+            return symm
+        
+        self.scaling_activation = torch.exp
+        self.scaling_inverse_activation = torch.log
+
+        self.covariance_activation = build_covariance_from_scaling_rotation
+
+        self.opacity_activation = torch.sigmoid
+        self.inverse_opacity_activation = torch.logit
+
+        self.rotation_activation = torch.nn.functional.normalize
+
+
+    def __init__(self, sh_degree : int, feat_mutable=True, affine=False):
+        self.active_sh_degree = 0
+        self.max_sh_degree = sh_degree  
+        self._xyz = torch.empty(0)
+        self._features_dc = torch.empty(0)
+        self._features_rest = torch.empty(0)
+        self._feats3D = torch.empty(0)
+        self._scaling = torch.empty(0)
+        self._rotation = torch.empty(0)
+        self._opacity = torch.empty(0)
+        self.max_radii2D = torch.empty(0)
+        self.xyz_gradient_accum = torch.empty(0)
+        self.denom = torch.empty(0)
+        self.optimizer = None
+        self.percent_dense = 0
+        self.spatial_lr_scale = 0
+        self.feat_mutable = feat_mutable
+        self.setup_functions()
+
+        self.pos_enc = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={"otype": "Frequency", "n_frequencies": 2},
+        )
+        self.dir_enc = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={
+                "otype": "SphericalHarmonics",
+                "degree": 3,
+            },
+        )
+
+        self.affine = affine
+        if affine:
+            self.appearance_model = tcnn.Network(
+                n_input_dims=self.pos_enc.n_output_dims + self.dir_enc.n_output_dims,
+                n_output_dims=12,
+                network_config={
+                    "otype": "FullyFusedMLP",
+                    "activation": "ReLU",
+                    "output_activation": "None",
+                    "n_neurons": 32,
+                    "n_hidden_layers": 2,
+                }
+            )
+        else:
+            self.appearance_model = None
+
+    def capture(self):
+        return (
+            self.active_sh_degree,
+            self._xyz,
+            self._features_dc,
+            self._features_rest,
+            self._feats3D,
+            self._scaling,
+            self._rotation,
+            self._opacity,
+            self.spatial_lr_scale,
+        )
+    
+    def restore(self, model_args, training_args):
+        (self.active_sh_degree, 
+        self._xyz, 
+        self._features_dc, 
+        self._features_rest,
+        self._feats3D,
+        self._scaling, 
+        self._rotation, 
+        self._opacity,
+        self.spatial_lr_scale,
+        ) = model_args
+        if training_args is not None:
+            self.training_setup(training_args)
+        
+    @property
+    def get_scaling(self):
+        return self.scaling_activation(self._scaling)
+        
+    @property
+    def get_rotation(self):
+        return self.rotation_activation(self._rotation)
+    
+    @property
+    def get_xyz(self):
+        return self._xyz
+
+    @property
+    def get_features(self):
+        features_dc = self._features_dc
+        features_rest = self._features_rest
+        return torch.cat((features_dc, features_rest), dim=1)
+    
+    @property
+    def get_3D_features(self):
+        return torch.softmax(self._feats3D, dim=-1)
+
+    @property
+    def get_opacity(self):
+        return self.opacity_activation(self._opacity)
+    
+    # def get_covariance(self, scaling_modifier = 1):
+    #     return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
+
+    def oneupSHdegree(self):
+        if self.active_sh_degree < self.max_sh_degree:
+            self.active_sh_degree += 1
+
+    def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float):
+        # self.spatial_lr_scale = 1
+        self.spatial_lr_scale = spatial_lr_scale
+        fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
+        fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
+        features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
+        features[:, :3, 0 ] = fused_color
+        features[:, 3:, 1:] = 0.0
+
+        if self.feat_mutable:
+            feats3D = torch.rand(fused_color.shape[0], 20).float().cuda()
+            self._feats3D = nn.Parameter(feats3D.requires_grad_(True))
+        else:
+            feats3D = torch.zeros(fused_color.shape[0], 20).float().cuda()
+            feats3D[:, 13] = 1
+            self._feats3D = feats3D
+
+        print("Number of points at initialization : ", fused_point_cloud.shape[0])
+
+        dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
+        scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
+        rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
+        rots[:, 0] = 1
+
+        opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
+
+        self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))
+        self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
+        self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
+        self._scaling = nn.Parameter(scales.requires_grad_(True))
+        self._rotation = nn.Parameter(rots.requires_grad_(True))
+        self._opacity = nn.Parameter(opacities.requires_grad_(True))
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def training_setup(self, training_args):
+        self.percent_dense = training_args.percent_dense
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+
+        # self.spatial_lr_scale /= 3
+
+        l = [
+            {'params': [self._xyz], 'lr': training_args.position_lr_init * 0.5, "name": "xyz"},
+            {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
+            {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
+            {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
+            {'params': [self._scaling], 'lr': training_args.scaling_lr * 0.5, "name": "scaling"},
+            {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"},
+        ]
+
+        if self.affine:
+            l.append({'params': [*self.appearance_model.parameters()], 'lr': 1e-3, "name": "appearance_model"})
+        
+        if self.feat_mutable:
+            l.append({'params': [self._feats3D], 'lr': 1e-2, "name": "feats3D"})
+
+        self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)
+        self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,
+                                                    lr_final=training_args.position_lr_final*self.spatial_lr_scale,
+                                                    lr_delay_mult=training_args.position_lr_delay_mult,
+                                                    max_steps=training_args.position_lr_max_steps)
+
+    def update_learning_rate(self, iteration):
+        ''' Learning rate scheduling per step '''
+        for param_group in self.optimizer.param_groups:
+            if param_group["name"] == "xyz":
+                lr = self.xyz_scheduler_args(iteration)
+                param_group['lr'] = lr
+                return lr
+
+    def construct_list_of_attributes(self):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        # All channels except the 3 DC
+        for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
+            l.append('f_dc_{}'.format(i))
+        for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):
+            l.append('f_rest_{}'.format(i))
+        for i in range(self._feats3D.shape[1]):
+            l.append('semantic_{}'.format(i))
+        l.append('opacity')
+        for i in range(self._scaling.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(self._rotation.shape[1]):
+            l.append('rot_{}'.format(i))
+        return l
+
+    def save_ply(self, path=None):
+        mkdir_p(os.path.dirname(path))
+
+        xyz = self.get_xyz.detach().cpu().numpy()
+        normals = np.zeros_like(xyz)
+        f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        feats3D = self._feats3D.detach().cpu().numpy()
+        opacities = self._opacity.detach().cpu().numpy()
+        scale = self.scaling_inverse_activation(self.get_scaling).detach().cpu().numpy()
+        rotation = self._rotation.detach().cpu().numpy()
+
+        dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyz, normals, f_dc, f_rest, feats3D, opacities, scale, rotation), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+        plydata = PlyData([el])
+        if path is not None:
+            plydata.write(path)
+        return plydata
+
+    def save_splat(self, ply_path, splat_path):
+        plydata = self.save_ply(ply_path)
+        vert = plydata["vertex"]
+        sorted_indices = np.argsort(
+            -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
+            / (1 + np.exp(-vert["opacity"]))
+        )
+        buffer = BytesIO()
+        for idx in sorted_indices:
+            v = plydata["vertex"][idx]
+            position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32)
+            scales = np.exp(
+                np.array(
+                    [v["scale_0"], v["scale_1"], v["scale_2"]],
+                    dtype=np.float32,
+                )
+            )
+            rot = np.array(
+                [v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]],
+                dtype=np.float32,
+            )
+            SH_C0 = 0.28209479177387814
+            color = np.array(
+                [
+                    0.5 + SH_C0 * v["f_dc_0"],
+                    0.5 + SH_C0 * v["f_dc_1"],
+                    0.5 + SH_C0 * v["f_dc_2"],
+                    1 / (1 + np.exp(-v["opacity"])),
+                ]
+            )
+            buffer.write(position.tobytes())
+            buffer.write(scales.tobytes())
+            buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
+            buffer.write(
+                ((rot / np.linalg.norm(rot)) * 128 + 128)
+                .clip(0, 255)
+                .astype(np.uint8)
+                .tobytes()
+            )
+        with open(splat_path, "wb") as f:
+            f.write(buffer.getvalue())
+
+    def save_semantic_pcd(self, path):
+        color_dict = {
+            0: np.array([128, 64, 128]),  # Road
+            1: np.array([244, 35, 232]),  # Sidewalk
+            2: np.array([70, 70, 70]),  # Building
+            3: np.array([102, 102, 156]),  # Wall
+            4: np.array([190, 153, 153]),  # Fence
+            5: np.array([153, 153, 153]),  # Pole
+            6: np.array([250, 170, 30]),  # Traffic Light
+            7: np.array([220, 220, 0]),  # Traffic Sign
+            8: np.array([107, 142, 35]),  # Vegetation
+            9: np.array([152, 251, 152]),  # Terrain
+            10: np.array([0, 0, 0]),  # Black (trainId 10)
+            11: np.array([70, 130, 180]),  # Sky
+            12: np.array([220, 20, 60]),  # Person
+            13: np.array([255, 0, 0]),  # Rider
+            14: np.array([0, 0, 142]),  # Car
+            15: np.array([0, 0, 70]),  # Truck
+            16: np.array([0, 60, 100]),  # Bus
+            17: np.array([0, 80, 100]),  # Train
+            18: np.array([0, 0, 230]),  # Motorcycle
+            19: np.array([119, 11, 32])  # Bicycle
+        }
+        semantic_idx = torch.argmax(self.get_3D_features, dim=-1, keepdim=True)
+        opacities = self.get_opacity[:, 0]
+        mask = ((semantic_idx != 10)[:, 0]) & ((semantic_idx != 8)[:, 0]) & (opacities > 0.2)
+
+        semantic_idx = semantic_idx[mask]
+        semantic_rgb = torch.zeros_like(semantic_idx).repeat(1, 3)
+        for idx in range(20):
+            rgb = torch.from_numpy(color_dict[idx]).to(semantic_rgb.device)[None, :]
+            semantic_rgb[(semantic_idx == idx)[:, 0], :] = rgb
+        semantic_rgb = semantic_rgb.float() / 255.0
+        pcd_xyz = self.get_xyz[mask]
+        smt_pcd = o3d.geometry.PointCloud()
+        smt_pcd.points = o3d.utility.Vector3dVector(pcd_xyz.detach().cpu().numpy())
+        smt_pcd.colors = o3d.utility.Vector3dVector(semantic_rgb.detach().cpu().numpy())
+        o3d.io.write_point_cloud(path, smt_pcd)
+
+    def save_vis_ply(self, path):
+        mkdir_p(os.path.dirname(path))
+        xyz = self.get_xyz.detach().cpu().numpy()
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(xyz)
+        colors = SH2RGB(self._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1)
+        pcd.colors = o3d.utility.Vector3dVector(colors)
+        o3d.io.write_point_cloud(path, pcd)
+
+    def reset_opacity(self):
+        opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
+        optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
+        self._opacity = optimizable_tensors["opacity"]
+
+    def load_ply(self, path):
+        plydata = PlyData.read(path)
+
+        xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                        np.asarray(plydata.elements[0]["y"]),
+                        np.asarray(plydata.elements[0]["z"])),  axis=1)
+        opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+        features_dc = np.zeros((xyz.shape[0], 3, 1))
+        features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+        features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
+        features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+        extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
+        assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3
+        features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
+        for idx, attr_name in enumerate(extra_f_names):
+            features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
+        # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
+        features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
+
+        scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+        scales = np.zeros((xyz.shape[0], len(scale_names)))
+        for idx, attr_name in enumerate(scale_names):
+            scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
+        rots = np.zeros((xyz.shape[0], len(rot_names)))
+        for idx, attr_name in enumerate(rot_names):
+            rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
+        self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
+        self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
+
+        self.active_sh_degree = self.max_sh_degree
+
+    def replace_tensor_to_optimizer(self, tensor, name):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group["name"] == name:
+                stored_state = self.optimizer.state.get(group['params'][0], None)
+                if stored_state is not None:
+                    stored_state["exp_avg"] = torch.zeros_like(tensor)
+                    stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
+                    del self.optimizer.state[group['params'][0]]
+                    group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                    self.optimizer.state[group['params'][0]] = stored_state
+                    optimizable_tensors[group["name"]] = group["params"][0]
+                else:
+                    group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                    optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def _prune_optimizer(self, mask):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group['name'] == 'appearance_model':
+                continue
+            stored_state = self.optimizer.state.get(group['params'][0], None)
+            if stored_state is not None:
+                stored_state["exp_avg"] = stored_state["exp_avg"][mask]
+                stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def prune_points(self, mask):
+        valid_points_mask = ~mask
+        optimizable_tensors = self._prune_optimizer(valid_points_mask)
+
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        if self.feat_mutable:
+            self._feats3D = optimizable_tensors["feats3D"]
+        else:
+            self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1))
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+
+        self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
+
+        self.denom = self.denom[valid_points_mask]
+        self.max_radii2D = self.max_radii2D[valid_points_mask]
+
+    def cat_tensors_to_optimizer(self, tensors_dict):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group['name'] not in tensors_dict:
+                continue
+            assert len(group["params"]) == 1
+            extension_tensor = tensors_dict[group["name"]]
+            stored_state = self.optimizer.state.get(group["params"][0], None)
+            if stored_state is not None:
+
+                stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
+                stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
+
+                del self.optimizer.state[group["params"][0]]
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                self.optimizer.state[group["params"][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+
+        return optimizable_tensors
+
+    def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation):
+        d = {"xyz": new_xyz,
+        "f_dc": new_features_dc,
+        "f_rest": new_features_rest,
+        "feats3D": new_feats3D,
+        "opacity": new_opacities,
+        "scaling" : new_scaling,
+        "rotation" : new_rotation}
+
+        optimizable_tensors = self.cat_tensors_to_optimizer(d)
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        if self.feat_mutable:
+            self._feats3D = optimizable_tensors["feats3D"]
+        else:
+            self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1))
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
+        n_init_points = self.get_xyz.shape[0]
+        # Extract points that satisfy the gradient condition
+        padded_grad = torch.zeros((n_init_points), device="cuda")
+        padded_grad[:grads.shape[0]] = grads.squeeze()
+        selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
+
+        stds = self.get_scaling[selected_pts_mask].repeat(N,1)
+        means =torch.zeros((stds.size(0), 3),device="cuda")
+        samples = torch.normal(mean=means, std=stds)
+        rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
+        new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
+        new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))
+        new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
+        new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
+        new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
+        new_feats3D = self._feats3D[selected_pts_mask].repeat(N,1)
+        new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacity, new_scaling, new_rotation)
+
+        prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
+        self.prune_points(prune_filter)
+
+    def densify_and_clone(self, grads, grad_threshold, scene_extent):
+        # Extract points that satisfy the gradient condition
+        selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
+        
+        new_xyz = self._xyz[selected_pts_mask]
+        new_features_dc = self._features_dc[selected_pts_mask]
+        new_features_rest = self._features_rest[selected_pts_mask]
+        new_feats3D = self._feats3D[selected_pts_mask]
+        new_opacities = self._opacity[selected_pts_mask]
+        new_scaling = self._scaling[selected_pts_mask]
+        new_rotation = self._rotation[selected_pts_mask]
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation)
+
+    def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size, cam_pos=None):
+        grads = self.xyz_gradient_accum / self.denom
+        grads[grads.isnan()] = 0.0
+
+        self.densify_and_clone(grads, max_grad, extent)
+        self.densify_and_split(grads, max_grad, extent)
+
+        prune_mask = (self.get_opacity < min_opacity).squeeze()
+        if max_screen_size:
+            big_points_vs = self.max_radii2D > max_screen_size
+            if cam_pos is not None:
+                # points_cam_dist = torch.abs(self.get_xyz[:, None, :] - cam_pos[None, ...])
+                # points_cam_nearest_idx = torch.argmin(torch.norm(points_cam_dist, dim=-1), dim=1)
+                # points_cam_dist = points_cam_dist[torch.arange(points_cam_dist.shape[0]), points_cam_nearest_idx, :]
+                # near_mask1 = (points_cam_dist[:, 1] < 5) & (points_cam_dist[:, 0] < 10) & (points_cam_dist[:, 2] < 10)
+                # big_points_ws1 = near_mask1 & (self.get_scaling.max(dim=1).values > 1.0)
+                # near_mask2 = (points_cam_dist[:, 1] < 10) & (points_cam_dist[:, 0] < 20) & (points_cam_dist[:, 2] < 20)
+                # big_points_ws2 = near_mask2 & (self.get_scaling.max(dim=1).values > 3.0)
+                # big_points_ws = (self.get_scaling.max(dim=1).values > 10.0) | big_points_ws1 | big_points_ws2 
+                big_points_ws = self.get_scaling.max(dim=1).values > 10
+                prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+            else:
+                big_points_ws = self.get_scaling.max(dim=1).values > 5
+                prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+        self.prune_points(prune_mask)
+
+        torch.cuda.empty_cache()
+    
+    def add_densification_stats_grad(self, tensor_grad, update_filter):
+        self.xyz_gradient_accum[update_filter] += torch.norm(tensor_grad[update_filter,:2], dim=-1, keepdim=True)
+        self.denom[update_filter] += 1
+        

code/sim/hugsim_env.egg-info/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.4
+Name: hugsim-env
+Version: 0.0.1
+Requires-Dist: gymnasium

code/sim/hugsim_env.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+pyproject.toml
+setup.py
+hugsim_env/__init__.py
+hugsim_env.egg-info/PKG-INFO
+hugsim_env.egg-info/SOURCES.txt
+hugsim_env.egg-info/dependency_links.txt
+hugsim_env.egg-info/requires.txt
+hugsim_env.egg-info/top_level.txt
+hugsim_env/envs/__init__.py
+hugsim_env/envs/hug_sim.py

code/sim/hugsim_env.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

code/sim/hugsim_env.egg-info/requires.txt

@@ -0,0 +1 @@
+gymnasium

code/sim/hugsim_env.egg-info/top_level.txt

@@ -0,0 +1 @@
+hugsim_env

code/sim/hugsim_env/__init__.py

@@ -0,0 +1,8 @@
+from gymnasium.envs.registration import register
+
+
+register(
+     id="hugsim_env/HUGSim-v0",
+     entry_point="hugsim_env.envs:HUGSimEnv",
+     max_episode_steps=400,
+)

code/sim/hugsim_env/envs/__init__.py

@@ -0,0 +1 @@
+from hugsim_env.envs.hug_sim import HUGSimEnv

code/sim/hugsim_env/envs/hug_sim.py

@@ -0,0 +1,333 @@
+import torch
+import numpy as np
+from copy import deepcopy
+import gymnasium
+from gymnasium import spaces
+from copy import deepcopy
+from sim.utils.sim_utils import create_cam, rt2pose, pose2rt, load_camera_cfg, dense_cam_poses
+from scipy.spatial.transform import Rotation as SCR
+from sim.utils.score_calculator import create_rectangle, bg_collision_det
+import os
+import pickle
+from sim.utils.plan import planner, UnifiedMap
+from omegaconf import OmegaConf
+import math
+from gaussian_renderer import GaussianModel
+from scene.obj_model import ObjModel
+from gaussian_renderer import render
+import open3d as o3d
+
+
+def fg_collision_det(ego_box, objs):
+    ego_x, ego_y, _, ego_w, ego_l, ego_h, ego_yaw = ego_box
+    ego_poly = create_rectangle(ego_x, ego_y, ego_w, ego_l, ego_yaw)
+    for obs in objs:
+        obs_x, obs_y, _, obs_w, obs_l, _, obs_yaw = obs
+        obs_poly = create_rectangle(
+            obs_x, obs_y, obs_w, obs_l, obs_yaw)
+        if ego_poly.intersects(obs_poly):
+            return True
+    return False
+
+class HUGSimEnv(gymnasium.Env):
+    def __init__(self, cfg, output):
+        super().__init__()
+        
+        plan_list = cfg.scenario.plan_list
+        for control_param in plan_list:
+            control_param[5] = os.path.join(cfg.base.realcar_path, control_param[5])
+
+        # read ground infos
+        with open(os.path.join(cfg.model_path, 'ground_param.pkl'), 'rb') as f:
+            #numpy.ndarray, float, list
+            cam_poses, cam_heights, commands = pickle.load(f)
+            cam_poses, commands = dense_cam_poses(cam_poses, commands)
+            self.ground_model = (cam_poses, cam_heights, commands)
+
+        if cfg.scenario.load_HD_map:
+            unified_map = UnifiedMap(cfg.base.HD_map.path, cfg.base.HD_map.version, cfg.scenario.scene_name)
+        else:
+            unified_map = None
+        
+        self.kinematic = OmegaConf.to_container(cfg.kinematic)
+        self.kinematic['min_steer'] = -math.radians(cfg.kinematic.min_steer)
+        self.kinematic['max_steer'] = math.radians(cfg.kinematic.max_steer)
+        self.kinematic['start_vr']= np.array(cfg.scenario.start_euler) / 180 * np.pi
+        self.kinematic['start_vab'] = np.array(cfg.scenario.start_ab)
+        self.kinematic['start_velo'] = cfg.scenario.start_velo
+        self.kinematic['start_steer'] = cfg.scenario.start_steer
+
+        self.gaussians = GaussianModel(cfg.model.sh_degree, affine=cfg.affine)
+
+        """
+        plan_list: a, b, height, yaw, v, model_path, controller, params
+        Yaw is based on ego car's orientation. 0 means same direction as ego. 
+        Right is positive and left is negative.
+        """
+        
+        (model_params, iteration) = torch.load(os.path.join(cfg.model_path, "scene.pth"), weights_only=False)
+        self.gaussians.restore(model_params, None)
+        
+        dynamic_gaussians = {}
+        if len(plan_list) == 0:
+            self.planner = None
+        else:
+            self.planner = planner(plan_list, scene_path=cfg.model_path, unified_map=unified_map, ground=self.ground_model, dt=cfg.kinematic.dt)
+            for plan_id in self.planner.ckpts.keys():
+                dynamic_gaussians[plan_id] = ObjModel(cfg.model.sh_degree, feat_mutable=False)
+                (model_params, iteration) = torch.load(self.planner.ckpts[plan_id], weights_only=False)
+                model_params = list(model_params)
+                dynamic_gaussians[plan_id].restore(model_params, None)
+            
+        semantic_idx = torch.argmax(self.gaussians.get_full_3D_features, dim=-1, keepdim=True)
+        ground_xyz = self.gaussians.get_full_xyz[(semantic_idx == 0)[:, 0]].detach().cpu().numpy()
+        scene_xyz = self.gaussians.get_full_xyz[((semantic_idx > 1) & (semantic_idx != 10))[:, 0]].detach().cpu().numpy()
+        ground_pcd = o3d.geometry.PointCloud()
+        ground_pcd.points = o3d.utility.Vector3dVector(ground_xyz.astype(float))
+        o3d.io.write_point_cloud(os.path.join(output, 'ground.ply'), ground_pcd)
+        scene_pcd = o3d.geometry.PointCloud()
+        scene_pcd.points = o3d.utility.Vector3dVector(scene_xyz.astype(float))
+        o3d.io.write_point_cloud(os.path.join(output, 'scene.ply'), scene_pcd)
+            
+        unicycles = {}
+
+        if cfg.scenario.load_HD_map and self.planner is not None:
+            self.planner.update_agent_route()
+        
+        self.cam_params, cam_align, self.cam_rect = load_camera_cfg(cfg.camera)
+        
+        self.ego_verts = np.array([[0.5, 0, 0.5], [0.5, 0, -0.5], [0.5, 1.0,  0.5], [0.5, 1.0, -0.5],
+                    [-0.5, 0, -0.5], [-0.5, 0, 0.5], [-0.5, 1.0, -0.5], [-0.5, 1.0, 0.5]])
+        self.whl = np.array([1.6, 1.5, 3.0])
+        self.ego_verts *= self.whl
+        self.data_type = cfg.data_type
+
+        self.action_space = spaces.Dict(
+            {
+                "steer_rate": spaces.Box(self.kinematic['min_steer'], self.kinematic['max_steer'], dtype=float),
+                "acc": spaces.Box(self.kinematic['min_acc'], self.kinematic['max_acc'], dtype=float)
+            }
+        )
+        self.observation_space = spaces.Dict(
+            {
+                'rgb': spaces.Dict({
+                    cam_name: spaces.Box(
+                        low=0, high=255, 
+                        shape=(params['intrinsic']['H'], params['intrinsic']['W'], 3), dtype=np.uint8
+                    ) for cam_name, params in self.cam_params.items()
+                }),
+                # 'semantic': spaces.Dict({
+                #     cam_name: spaces.Box(
+                #         low=0, high=50, 
+                #         shape=(params['intrinsic']['H'], params['intrinsic']['W']), dtype=np.uint8
+                #     ) for cam_name, params in self.cam_params.items()
+                # }),
+                # 'depth': spaces.Dict({
+                #     cam_name: spaces.Box(
+                #         low=0, high=1000, 
+                #         shape=(params['intrinsic']['H'], params['intrinsic']['W']), dtype=np.float32
+                #     ) for cam_name, params in self.cam_params.items()
+                # }),
+            }
+        )
+        self.fric = self.kinematic['fric']
+
+        self.start_vr = self.kinematic['start_vr']
+        self.start_vab = self.kinematic['start_vab']
+        self.start_velo = self.kinematic['start_velo']
+        self.vr = deepcopy(self.kinematic['start_vr'])
+        self.vab = deepcopy(self.kinematic['start_vab'])
+        self.velo = deepcopy(self.kinematic['start_velo'])
+        self.steer = deepcopy(self.kinematic['start_steer'])
+        self.dt = self.kinematic['dt']
+
+        bg_color = [1, 1, 1] if cfg.model.white_background else [0, 0, 0]
+        self.render_fn = render
+        self.render_kwargs = {
+            "pc": self.gaussians,
+            "bg_color": torch.tensor(bg_color, dtype=torch.float32, device="cuda"),
+            "dynamic_gaussians": dynamic_gaussians,
+            "unicycles": unicycles
+        }
+        gaussians = self.gaussians
+        semantic_idx = torch.argmax(gaussians.get_3D_features, dim=-1, keepdim=True)
+        opacities = gaussians.get_opacity[:, 0]
+        mask = ((semantic_idx > 1) & (semantic_idx != 10))[:, 0] & (opacities > 0.8)
+        self.points = gaussians.get_xyz[mask]
+
+        self.last_accel = 0
+        self.last_steer_rate = 0
+
+        self.timestamp = 0
+    
+    def ground_height(self, u, v):
+        cam_poses, cam_height, _ = self.ground_model
+        cam_dist = np.sqrt(
+            (cam_poses[:, 0, 3] - u)**2 + (cam_poses[:, 2, 3] - v)**2
+        )
+        nearest_cam_idx = np.argmin(cam_dist, axis=0)
+        nearest_c2w = cam_poses[nearest_cam_idx]
+
+        nearest_w2c = np.linalg.inv(nearest_c2w)
+        uhv_local = nearest_w2c[:3, :3] @ np.array([u, 0, v]) + nearest_w2c[:3, 3]
+        uhv_local[1] = 0
+        uhv_world = nearest_c2w[:3, :3] @ uhv_local + nearest_c2w[:3, 3]
+        
+        return uhv_world[1]
+    
+    @property
+    def route_completion(self):
+        cam_poses, _, _ = self.ground_model
+        cam_dist = np.sqrt(
+            (cam_poses[:, 0, 3] - self.vab[0])**2 + (cam_poses[:, 2, 3] - self.vab[1])**2
+        )
+        nearest_cam_idx = np.argmin(cam_dist, axis=0)
+        return (nearest_cam_idx + 1) / (cam_poses.shape[0] * 0.9), cam_dist[nearest_cam_idx]
+        
+
+    @property
+    def vt(self):
+        vt = np.zeros(3)
+        vt[[0, 2]] = self.vab
+        vt[1] = self.ground_height(self.vab[0], self.vab[1])
+        return vt
+    
+    @property
+    def ego(self):
+        return rt2pose(self.vr, self.vt)
+    
+    @property
+    def ego_state(self):
+        return torch.tensor([self.vab[0], self.vab[1], self.vr[1], self.velo])
+    
+    @property
+    def ego_box(self):
+        return [self.vt[2], -self.vt[0], -self.vt[1], self.whl[0], self.whl[2], self.whl[1], -self.vr[1]]
+
+    @property
+    def objs_list(self):
+        obj_boxes = []
+        objs = self.render_kwargs['planning'][0]
+        for obj_id, obj_b2w in objs.items():
+            yaw = SCR.from_matrix(obj_b2w[:3, :3].detach().cpu().numpy()).as_euler('YXZ')[0]
+            # X, Y, Z in IMU, w, l, h
+            wlh = self.planner.wlhs[obj_id]
+            obj_boxes.append([obj_b2w[2, 3].item(), -obj_b2w[0, 3].item(), -obj_b2w[1, 3].item(), wlh[0], wlh[1], wlh[2], -yaw-0.5*np.pi])
+        return obj_boxes
+
+    def _get_obs(self):
+        rgbs, semantics, depths = {}, {}, {}
+        v2front = self.cam_params['CAM_FRONT']["v2c"]
+        for cam_name, params in self.cam_params.items():
+            intrinsic, v2c = params['intrinsic'], params['v2c']
+            c2front = v2front @ np.linalg.inv(v2c) @ self.cam_rect
+            c2w = self.ego @ c2front
+            viewpoint = create_cam(intrinsic, c2w)
+            with torch.no_grad():
+                render_pkg = self.render_fn(viewpoint=viewpoint, prev_viewpoint=None, **self.render_kwargs)
+            rgb = (torch.permute(render_pkg['render'].clamp(0, 1), (1,2,0)).detach().cpu().numpy() * 255).astype(np.uint8)
+            smt = torch.argmax(render_pkg['feats'], dim=0).detach().cpu().numpy().astype(np.uint8)
+            depth = render_pkg['depth'][0].detach().cpu().numpy()
+            if (self.data_type == 'waymo' or self.data_type == 'kitti360') and 'BACK' in cam_name:
+                rgbs[cam_name] = np.zeros_like(rgb)
+                semantics[cam_name] = np.zeros_like(smt)
+                depths[cam_name] = np.zeros_like(depth)
+            else:
+                rgbs[cam_name] = rgb
+                semantics[cam_name] = smt
+                depths[cam_name] = depth
+
+        return {
+                'rgb': rgbs, 
+                # 'semantic': semantics,
+                # 'depth': depths,
+                }
+    
+    def _get_info(self):
+        wego_r, wego_t = pose2rt(self.ego)
+        cam_poses, _, commands = self.ground_model
+        dist = np.sum((cam_poses[:, :3, 3] - self.vt) ** 2, axis=-1)
+        nearest_cam_idx = np.argmin(dist)
+        command = commands[nearest_cam_idx]
+        return {
+            'ego_pos'  : wego_t.tolist(),
+            'ego_rot'  : wego_r.tolist(),
+            'ego_velo' : self.velo,
+            'ego_steer': self.steer,
+            'accelerate': self.last_accel,
+            'steer_rate': self.last_steer_rate,
+            'timestamp': self.timestamp,
+            'command': command,
+            'ego_box': self.ego_box,
+            'obj_boxes': self.objs_list,
+            'cam_params': self.cam_params,
+            # 'ego_verts': verts,
+        }
+    
+    def reset(self, seed=None, options=None):
+        self.vr = deepcopy(self.start_vr)
+        self.vab = deepcopy(self.start_vab)
+        self.velo = deepcopy(self.start_velo)
+        self.timestamp = 0
+
+        if self.planner is not None:
+            self.render_kwargs['planning'] = self.planner.plan_traj(self.timestamp, self.ego_state)
+        else:
+            self.render_kwargs['planning'] = [{},  {}]
+            
+        observation = self._get_obs()
+        info = self._get_info()
+
+        return observation, info
+    
+    def step(self, action):
+        self.timestamp += self.dt
+        if self.planner is not None:
+            self.render_kwargs['planning'] = self.planner.plan_traj(self.timestamp, self.ego_state)
+        else:
+            self.render_kwargs['planning'] = [{},  {}]
+        steer_rate, acc = action['steer_rate'], action['acc']
+        self.last_steer_rate, self.last_accel = steer_rate, acc
+        L = self.kinematic['Lr'] + self.kinematic['Lf']
+        self.velo += acc * self.dt
+        self.steer += steer_rate * self.dt
+        theta = self.vr[1]
+        # print(theta / np.pi * 180, self.steer / np.pi * 180)
+        self.vab[0] = self.vab[0] + self.velo * np.sin(theta) * self.dt
+        self.vab[1] = self.vab[1] + self.velo * np.cos(theta) * self.dt
+        self.vr[1] = theta + self.velo * np.tan(self.steer) / L * self.dt
+
+        terminated = False
+        reward = 0
+        verts = (self.ego[:3, :3] @ self.ego_verts.T).T + self.ego[:3, 3]
+        verts = torch.from_numpy(verts.astype(np.float32)).cuda()
+        
+        bg_collision = bg_collision_det(self.points, verts)
+        if bg_collision:
+            terminated = True
+            print('Collision with background')
+            reward = -100
+
+        fg_collision = fg_collision_det(self.ego_box, self.objs_list)
+        if fg_collision:
+            terminated = True
+            print('Collision with foreground')
+            reward = -100
+
+        rc, dist = self.route_completion
+        if dist > 10:
+            terminated=True
+            print('Far from preset trajectory')
+            reward = -50
+            
+        if rc >= 1:
+            terminated = True
+            print('Complete')
+            reward = 1000
+
+        observation = self._get_obs()
+        info = self._get_info()
+        info['rc'] = rc
+        info['collision'] = bg_collision or fg_collision
+        
+        return observation, reward, terminated, False, info

code/sim/ilqr/lqr.py

@@ -0,0 +1,55 @@
+from sim.ilqr.lqr_solver import ILQRSolverParameters, ILQRWarmStartParameters, ILQRSolver
+import numpy as np
+
+solver_params = ILQRSolverParameters(
+    discretization_time=0.5,
+    state_cost_diagonal_entries=[1.0, 1.0, 10.0, 0.0, 0.0],
+    input_cost_diagonal_entries=[1.0, 10.0],
+    state_trust_region_entries=[1.0] * 5,
+    input_trust_region_entries=[1.0] * 2,
+    max_ilqr_iterations=100,
+    convergence_threshold=1e-6,
+    max_solve_time=0.05,
+    max_acceleration=3.0,
+    max_steering_angle=np.pi / 3.0,
+    max_steering_angle_rate=0.4,
+    min_velocity_linearization=0.01,
+    wheelbase=2.7
+)
+
+warm_start_params = ILQRWarmStartParameters(
+    k_velocity_error_feedback=0.5,
+    k_steering_angle_error_feedback=0.05,
+    lookahead_distance_lateral_error=15.0,
+    k_lateral_error=0.1,
+    jerk_penalty_warm_start_fit=1e-4,
+    curvature_rate_penalty_warm_start_fit=1e-2,
+)
+
+lqr = ILQRSolver(solver_params=solver_params, warm_start_params=warm_start_params)
+
+def plan2control(plan_traj, init_state):
+    current_state = init_state
+    solutions = lqr.solve(current_state, plan_traj)
+    optimal_inputs = solutions[-1].input_trajectory
+    accel_cmd = optimal_inputs[0, 0]
+    steering_rate_cmd = optimal_inputs[0, 1]
+    return accel_cmd, steering_rate_cmd
+
+if __name__ == '__main__':
+    # plan_traj = np.zeros((6,5))
+    # plan_traj[:, 0] = 1
+    # plan_traj[:, 1] = np.ones(6)
+    # plan_traj = np.cumsum(plan_traj, axis=0)
+    # print(plan_traj)
+    plan_traj = np.array([[-0.18724936,  2.29100776,  0.,          0.,          0.,        ],
+                        [-0.29260731,  2.2971828 ,  0.,          0.,          0.        ],
+                        [-0.46831554,  2.55596018,  0.,          0.,          0.        ],
+                        [-0.5859955 ,  2.73183298,  0.,          0.,          0.        ],
+                        [-0.62684   ,  2.84659386,  0.,          0.,          0.        ],
+                        [-0.67761713,  2.80647802,  0.,          0.,          0.        ]])
+    plan_traj = plan_traj[:, [1,0,2,3,4]]
+    init_state = np.array([0.00000000e+00, 3.46944695e-17, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00])
+    print(plan_traj.shape, init_state.shape)
+    acc, steer = plan2control(plan_traj, init_state)
+    print(acc, steer)

code/sim/ilqr/lqr_solver.py

@@ -0,0 +1,689 @@
+"""
+This provides an implementation of the iterative linear quadratic regulator (iLQR) algorithm for trajectory tracking.
+It is specialized to the case with a discrete-time kinematic bicycle model and a quadratic trajectory tracking cost.
+
+Original (Nonlinear) Discrete Time System:
+    z_k = [x_k, y_k, theta_k, v_k, delta_k]
+    u_k = [a_k, phi_k]
+
+    x_{k+1}     = x_k     + v_k * cos(theta_k) * dt
+    y_{k+1}     = y_k     + v_k * sin(theta_k) * dt
+    theta_{k+1} = theta_k + v_k * tan(delta_k) / L * dt
+    v_{k+1}     = v_k     + a_k * dt
+    delta_{k+1} = delta_k + phi_k * dt
+
+    where (x_k, y_k, theta_k) is the pose at timestep k with time discretization dt,
+    v_k and a_k are velocity and acceleration,
+    delta_k and phi_k are steering angle and steering angle rate,
+    and L is the vehicle wheelbase.
+
+Quadratic Tracking Cost:
+    J = sum_{k=0}^{N-1} ||u_k||_2^{R_k} +
+        sum_{k=0}^N ||z_k - z_{ref,k}||_2^{Q_k}
+For simplicity, we opt to use constant input cost matrices R_k = R and constant state cost matrices Q_k = Q.
+
+There are multiple improvements that can be done for this implementation, but omitted for simplicity of the code.
+Some of these include:
+  * Handle constraints directly in the optimization (e.g. log-barrier / penalty method with quadratic cost estimate).
+  * Line search in the input policy update (feedforward term) to determine a good gradient step size.
+
+References Used: https://people.eecs.berkeley.edu/~pabbeel/cs287-fa19/slides/Lec5-LQR.pdf and
+                 https://www.cs.cmu.edu/~rsalakhu/10703/Lectures/Lecture_trajectoryoptimization.pdf
+"""
+
+import time
+from dataclasses import dataclass, fields
+from typing import List, Optional, Tuple
+
+import numpy as np
+import numpy.typing as npt
+
+# from nuplan.common.actor_state.vehicle_parameters import get_pacifica_parameters
+# from nuplan.common.geometry.compute import principal_value
+# from nuplan.planning.simulation.controller.tracker.tracker_utils import (
+#     complete_kinematic_state_and_inputs_from_poses,
+#     compute_steering_angle_feedback,
+# )
+from sim.ilqr.utils import principal_value, complete_kinematic_state_and_inputs_from_poses, compute_steering_angle_feedback
+
+DoubleMatrix = npt.NDArray[np.float64]
+
+
+@dataclass(frozen=True)
+class ILQRSolverParameters:
+    """Parameters related to the solver implementation."""
+
+    discretization_time: float  # [s] Time discretization used for integration.
+
+    # Cost weights for state [x, y, heading, velocity, steering angle] and input variables [acceleration, steering rate].
+    state_cost_diagonal_entries: List[float]
+    input_cost_diagonal_entries: List[float]
+
+    # Trust region cost weights for state and input variables.  Helps keep linearization error per update step bounded.
+    state_trust_region_entries: List[float]
+    input_trust_region_entries: List[float]
+
+    # Parameters related to solver runtime / solution sub-optimality.
+    max_ilqr_iterations: int  # Maximum number of iterations to run iLQR before timeout.
+    convergence_threshold: float  # Threshold for delta inputs below which we can terminate iLQR early.
+    max_solve_time: Optional[
+        float
+    ]  # [s] If defined, sets a maximum time to run a solve call of iLQR before terminating.
+
+    # Constraints for underlying dynamics model.
+    max_acceleration: float  # [m/s^2] Absolute value threshold on acceleration input.
+    max_steering_angle: float  # [rad] Absolute value threshold on steering angle state.
+    max_steering_angle_rate: float  # [rad/s] Absolute value threshold on steering rate input.
+
+    # Parameters for dynamics / linearization.
+    min_velocity_linearization: float  # [m/s] Absolute value threshold below which linearization velocity is modified.
+    wheelbase: float  # [m] Wheelbase length parameter for the vehicle.
+
+    def __post_init__(self) -> None:
+        """Ensure entries lie in expected bounds and initialize wheelbase."""
+        for entry in [
+            "discretization_time",
+            "max_ilqr_iterations",
+            "convergence_threshold",
+            "max_acceleration",
+            "max_steering_angle",
+            "max_steering_angle_rate",
+            "min_velocity_linearization",
+            "wheelbase",
+        ]:
+            assert getattr(self, entry) > 0.0, f"Field {entry} should be positive."
+
+        assert self.max_steering_angle < np.pi / 2.0, "Max steering angle should be less than 90 degrees."
+
+        if isinstance(self.max_solve_time, float):
+            assert self.max_solve_time > 0.0, "The specified max solve time should be positive."
+
+        assert np.all([x >= 0 for x in self.state_cost_diagonal_entries]), "Q matrix must be positive semidefinite."
+        assert np.all([x > 0 for x in self.input_cost_diagonal_entries]), "R matrix must be positive definite."
+
+        assert np.all(
+            [x > 0 for x in self.state_trust_region_entries]
+        ), "State trust region cost matrix must be positive definite."
+        assert np.all(
+            [x > 0 for x in self.input_trust_region_entries]
+        ), "Input trust region cost matrix must be positive definite."
+
+
+@dataclass(frozen=True)
+class ILQRWarmStartParameters:
+    """Parameters related to generating a warm start trajectory for iLQR."""
+
+    k_velocity_error_feedback: float  # Gain for initial velocity error for warm start acceleration.
+    k_steering_angle_error_feedback: float  # Gain for initial steering angle error for warm start steering rate.
+    lookahead_distance_lateral_error: float  # [m] Distance ahead for which we estimate lateral error.
+    k_lateral_error: float  # Gain for lateral error to compute steering angle feedback.
+    jerk_penalty_warm_start_fit: float  # Penalty for jerk in velocity profile estimation.
+    curvature_rate_penalty_warm_start_fit: float  # Penalty for curvature rate in curvature profile estimation.
+
+    def __post_init__(self) -> None:
+        """Ensure entries lie in expected bounds."""
+        for entry in [
+            "k_velocity_error_feedback",
+            "k_steering_angle_error_feedback",
+            "lookahead_distance_lateral_error",
+            "k_lateral_error",
+            "jerk_penalty_warm_start_fit",
+            "curvature_rate_penalty_warm_start_fit",
+        ]:
+            assert getattr(self, entry) > 0.0, f"Field {entry} should be positive."
+
+
+@dataclass(frozen=True)
+class ILQRIterate:
+    """Contains state, input, and associated Jacobian trajectories needed to perform an update step of iLQR."""
+
+    state_trajectory: DoubleMatrix
+    input_trajectory: DoubleMatrix
+    state_jacobian_trajectory: DoubleMatrix
+    input_jacobian_trajectory: DoubleMatrix
+
+    def __post_init__(self) -> None:
+        """Check consistency of dimension across trajectory elements."""
+        assert len(self.state_trajectory.shape) == 2, "Expect state trajectory to be a 2D matrix."
+        state_trajectory_length, state_dim = self.state_trajectory.shape
+
+        assert len(self.input_trajectory.shape) == 2, "Expect input trajectory to be a 2D matrix."
+        input_trajectory_length, input_dim = self.input_trajectory.shape
+
+        assert (
+            input_trajectory_length == state_trajectory_length - 1
+        ), "State trajectory should be 1 longer than the input trajectory."
+        assert self.state_jacobian_trajectory.shape == (input_trajectory_length, state_dim, state_dim)
+        assert self.input_jacobian_trajectory.shape == (input_trajectory_length, state_dim, input_dim)
+
+        for field in fields(self):
+            # Make sure that we have no nan entries in our trajectory rollout prior to operating on this.
+            assert ~np.any(np.isnan(getattr(self, field.name))), f"{field.name} has unexpected nan values."
+
+
+@dataclass(frozen=True)
+class ILQRInputPolicy:
+    """Contains parameters for the perturbation input policy computed after performing LQR."""
+
+    state_feedback_matrices: DoubleMatrix
+    feedforward_inputs: DoubleMatrix
+
+    def __post__init__(self) -> None:
+        """Check shape of policy parameters."""
+        assert (
+            len(self.state_feedback_matrices.shape) == 3
+        ), "Expected state_feedback_matrices to have shape (n_horizon, n_inputs, n_states)"
+
+        assert (
+            len(self.feedforward_inputs.shape) == 2
+        ), "Expected feedforward inputs to have shape (n_horizon, n_inputs)."
+
+        assert (
+            self.feedforward_inputs.shape == self.state_feedback_matrices.shape[:2]
+        ), "Inconsistent horizon or input dimension between feedforward inputs and state feedback matrices."
+
+        for field in fields(self):
+            # Make sure that we have no nan entries in our policy parameters prior to using them.
+            assert ~np.any(np.isnan(getattr(self, field.name))), f"{field.name} has unexpected nan values."
+
+
+@dataclass(frozen=True)
+class ILQRSolution:
+    """Contains the iLQR solution with associated cost for consumption by the solver's client."""
+
+    state_trajectory: DoubleMatrix
+    input_trajectory: DoubleMatrix
+
+    tracking_cost: float
+
+    def __post_init__(self) -> None:
+        """Check consistency of dimension across trajectory elements and nonnegative cost."""
+        assert len(self.state_trajectory.shape) == 2, "Expect state trajectory to be a 2D matrix."
+        state_trajectory_length, _ = self.state_trajectory.shape
+
+        assert len(self.input_trajectory.shape) == 2, "Expect input trajectory to be a 2D matrix."
+        input_trajectory_length, _ = self.input_trajectory.shape
+
+        assert (
+            input_trajectory_length == state_trajectory_length - 1
+        ), "State trajectory should be 1 longer than the input trajectory."
+
+        assert self.tracking_cost >= 0.0, "Expect the tracking cost to be nonnegative."
+
+
+class ILQRSolver:
+    """iLQR solver implementation, see module docstring for details."""
+
+    def __init__(
+        self,
+        solver_params: ILQRSolverParameters,
+        warm_start_params: ILQRWarmStartParameters,
+    ) -> None:
+        """
+        Initialize solver parameters.
+        :param solver_params: Contains solver parameters for iLQR.
+        :param warm_start_params: Contains warm start parameters for iLQR.
+        """
+        self._solver_params = solver_params
+        self._warm_start_params = warm_start_params
+
+        self._n_states = 5  # state dimension
+        self._n_inputs = 2  # input dimension
+
+        state_cost_diagonal_entries = self._solver_params.state_cost_diagonal_entries
+        assert (
+            len(state_cost_diagonal_entries) == self._n_states
+        ), f"State cost matrix should have diagonal length {self._n_states}."
+        self._state_cost_matrix: DoubleMatrix = np.diag(state_cost_diagonal_entries)
+
+        input_cost_diagonal_entries = self._solver_params.input_cost_diagonal_entries
+        assert (
+            len(input_cost_diagonal_entries) == self._n_inputs
+        ), f"Input cost matrix should have diagonal length {self._n_inputs}."
+        self._input_cost_matrix: DoubleMatrix = np.diag(input_cost_diagonal_entries)
+
+        state_trust_region_entries = self._solver_params.state_trust_region_entries
+        assert (
+            len(state_trust_region_entries) == self._n_states
+        ), f"State trust region cost matrix should have diagonal length {self._n_states}."
+        self._state_trust_region_cost_matrix: DoubleMatrix = np.diag(state_trust_region_entries)
+
+        input_trust_region_entries = self._solver_params.input_trust_region_entries
+        assert (
+            len(input_trust_region_entries) == self._n_inputs
+        ), f"Input trust region cost matrix should have diagonal length {self._n_inputs}."
+        self._input_trust_region_cost_matrix: DoubleMatrix = np.diag(input_trust_region_entries)
+
+        max_acceleration = self._solver_params.max_acceleration
+        max_steering_angle_rate = self._solver_params.max_steering_angle_rate
+
+        # Define input clip limits once to avoid recomputation in _clip_inputs.
+        self._input_clip_min = (-max_acceleration, -max_steering_angle_rate)
+        self._input_clip_max = (max_acceleration, max_steering_angle_rate)
+
+    def solve(self, current_state: DoubleMatrix, reference_trajectory: DoubleMatrix) -> List[ILQRSolution]:
+        """
+        Run the main iLQR loop used to try to find (locally) optimal inputs to track the reference trajectory.
+        :param current_state: The initial state from which we apply inputs, z_0.
+        :param reference_trajectory: The state reference we'd like to track, inclusive of the initial timestep,
+                                     z_{r,k} for k in {0, ..., N}.
+        :return: A list of solution iterates after running the iLQR algorithm where the index is the iteration number.
+        """
+        # Check that state parameter has the right shape.
+        assert current_state.shape == (self._n_states,), "Incorrect state shape."
+
+        # Check that reference trajectory parameter has the right shape.
+        assert len(reference_trajectory.shape) == 2, "Reference trajectory should be a 2D matrix."
+        reference_trajectory_length, reference_trajectory_state_dimension = reference_trajectory.shape
+        assert reference_trajectory_length > 1, "The reference trajectory should be at least two timesteps long."
+        assert (
+            reference_trajectory_state_dimension == self._n_states
+        ), "The reference trajectory should have a matching state dimension."
+
+        # List of ILQRSolution results where the index corresponds to the iteration of iLQR.
+        solution_list: List[ILQRSolution] = []
+
+        # Get warm start input and state trajectory, as well as associated Jacobians.
+        current_iterate = self._input_warm_start(current_state, reference_trajectory)
+
+        # Main iLQR Loop.
+        solve_start_time = time.perf_counter()
+        for _ in range(self._solver_params.max_ilqr_iterations):
+            # Determine the cost and store the associated solution object.
+            tracking_cost = self._compute_tracking_cost(
+                iterate=current_iterate,
+                reference_trajectory=reference_trajectory,
+            )
+            solution_list.append(
+                ILQRSolution(
+                    input_trajectory=current_iterate.input_trajectory,
+                    state_trajectory=current_iterate.state_trajectory,
+                    tracking_cost=tracking_cost,
+                )
+            )
+
+            # Determine the LQR optimal perturbations to apply.
+            lqr_input_policy = self._run_lqr_backward_recursion(
+                current_iterate=current_iterate,
+                reference_trajectory=reference_trajectory,
+            )
+
+            # Apply the optimal perturbations to generate the next input trajectory iterate.
+            input_trajectory_next = self._update_inputs_with_policy(
+                current_iterate=current_iterate,
+                lqr_input_policy=lqr_input_policy,
+            )
+
+            # Check for convergence/timeout and terminate early if so.
+            # Else update the input_trajectory iterate and continue.
+            input_trajectory_norm_difference = np.linalg.norm(input_trajectory_next - current_iterate.input_trajectory)
+
+            current_iterate = self._run_forward_dynamics(current_state, input_trajectory_next)
+
+            if input_trajectory_norm_difference < self._solver_params.convergence_threshold:
+                break
+
+            elapsed_time = time.perf_counter() - solve_start_time
+            if (
+                isinstance(self._solver_params.max_solve_time, float)
+                and elapsed_time >= self._solver_params.max_solve_time
+            ):
+                break
+
+        # Store the final iterate in the solution_dict.
+        tracking_cost = self._compute_tracking_cost(
+            iterate=current_iterate,
+            reference_trajectory=reference_trajectory,
+        )
+        solution_list.append(
+            ILQRSolution(
+                input_trajectory=current_iterate.input_trajectory,
+                state_trajectory=current_iterate.state_trajectory,
+                tracking_cost=tracking_cost,
+            )
+        )
+
+        return solution_list
+
+    ####################################################################################################################
+    # Helper methods.
+    ####################################################################################################################
+
+    def _compute_tracking_cost(self, iterate: ILQRIterate, reference_trajectory: DoubleMatrix) -> float:
+        """
+        Compute the trajectory tracking cost given a candidate solution.
+        :param iterate: Contains the candidate state and input trajectory to evaluate.
+        :param reference_trajectory: The desired state reference trajectory with same length as state_trajectory.
+        :return: The tracking cost of the candidate state/input trajectory.
+        """
+        input_trajectory = iterate.input_trajectory
+        state_trajectory = iterate.state_trajectory
+
+        assert len(state_trajectory) == len(
+            reference_trajectory
+        ), "The state and reference trajectory should have the same length."
+
+        error_state_trajectory = state_trajectory - reference_trajectory
+        error_state_trajectory[:, 2] = principal_value(error_state_trajectory[:, 2])
+
+        cost = np.sum([u.T @ self._input_cost_matrix @ u for u in input_trajectory]) + np.sum(
+            [e.T @ self._state_cost_matrix @ e for e in error_state_trajectory]
+        )
+
+        return float(cost)
+
+    def _clip_inputs(self, inputs: DoubleMatrix) -> DoubleMatrix:
+        """
+        Used to clip control inputs within constraints.
+        :param: inputs: The control inputs with shape (self._n_inputs,) to clip.
+        :return: Clipped version of the control inputs, unmodified if already within constraints.
+        """
+        assert inputs.shape == (self._n_inputs,), f"The inputs should be a 1D vector with {self._n_inputs} elements."
+
+        return np.clip(inputs, self._input_clip_min, self._input_clip_max)  # type: ignore
+
+    def _clip_steering_angle(self, steering_angle: float) -> float:
+        """
+        Used to clip the steering angle state within bounds.
+        :param steering_angle: [rad] A steering angle (scalar) to clip.
+        :return: [rad] The clipped steering angle.
+        """
+        steering_angle_sign = 1.0 if steering_angle >= 0 else -1.0
+        steering_angle = steering_angle_sign * min(abs(steering_angle), self._solver_params.max_steering_angle)
+        return steering_angle
+
+    def _input_warm_start(self, current_state: DoubleMatrix, reference_trajectory: DoubleMatrix) -> ILQRIterate:
+        """
+        Given a reference trajectory, we generate the warm start (initial guess) by inferring the inputs applied based
+        on poses in the reference trajectory.
+        :param current_state: The initial state from which we apply inputs.
+        :param reference_trajectory: The reference trajectory we are trying to follow.
+        :return: The warm start iterate from which to start iLQR.
+        """
+        reference_states_completed, reference_inputs_completed = complete_kinematic_state_and_inputs_from_poses(
+            discretization_time=self._solver_params.discretization_time,
+            wheel_base=self._solver_params.wheelbase,
+            poses=reference_trajectory[:, :3],
+            jerk_penalty=self._warm_start_params.jerk_penalty_warm_start_fit,
+            curvature_rate_penalty=self._warm_start_params.curvature_rate_penalty_warm_start_fit,
+        )
+
+        # We could just stop here and apply reference_inputs_completed (assuming it satisfies constraints).
+        # This could work if current_state = reference_states_completed[0,:] - i.e. no initial tracking error.
+        # We add feedback input terms for the first control input only to account for nonzero initial tracking error.
+        _, _, _, velocity_current, steering_angle_current = current_state
+        _, _, _, velocity_reference, steering_angle_reference = reference_states_completed[0, :]
+
+        acceleration_feedback = -self._warm_start_params.k_velocity_error_feedback * (
+            velocity_current - velocity_reference
+        )
+
+        steering_angle_feedback = compute_steering_angle_feedback(
+            pose_reference=current_state[:3],
+            pose_current=reference_states_completed[0, :3],
+            lookahead_distance=self._warm_start_params.lookahead_distance_lateral_error,
+            k_lateral_error=self._warm_start_params.k_lateral_error,
+        )
+        steering_angle_desired = steering_angle_feedback + steering_angle_reference
+        steering_rate_feedback = -self._warm_start_params.k_steering_angle_error_feedback * (
+            steering_angle_current - steering_angle_desired
+        )
+
+        reference_inputs_completed[0, 0] += acceleration_feedback
+        reference_inputs_completed[0, 1] += steering_rate_feedback
+
+        # We rerun dynamics with constraints applied to make sure we have a feasible warm start for iLQR.
+        return self._run_forward_dynamics(current_state, reference_inputs_completed)
+
+    ####################################################################################################################
+    # Dynamics and Jacobian.
+    ####################################################################################################################
+
+    def _run_forward_dynamics(self, current_state: DoubleMatrix, input_trajectory: DoubleMatrix) -> ILQRIterate:
+        """
+        Compute states and corresponding state/input Jacobian matrices using forward dynamics.
+        We additionally return the input since the dynamics may modify the input to ensure constraint satisfaction.
+        :param current_state: The initial state from which we apply inputs.  Must be feasible given constraints.
+        :param input_trajectory: The input trajectory applied to the model.  May be modified to ensure feasibility.
+        :return: A feasible iterate after applying dynamics with state/input trajectories and Jacobian matrices.
+        """
+        # Store rollout as a set of numpy arrays, initialized as np.nan to ensure we correctly fill them in.
+        # The state trajectory includes the current_state, z_0, and is 1 element longer than the other arrays.
+        # The final_input_trajectory captures the applied input for the dynamics model satisfying constraints.
+        N = len(input_trajectory)
+        state_trajectory = np.nan * np.ones((N + 1, self._n_states), dtype=np.float64)
+        final_input_trajectory = np.nan * np.ones_like(input_trajectory, dtype=np.float64)
+        state_jacobian_trajectory = np.nan * np.ones((N, self._n_states, self._n_states), dtype=np.float64)
+        final_input_jacobian_trajectory = np.nan * np.ones((N, self._n_states, self._n_inputs), dtype=np.float64)
+
+        state_trajectory[0] = current_state
+
+        for idx_u, u in enumerate(input_trajectory):
+            state_next, final_input, state_jacobian, final_input_jacobian = self._dynamics_and_jacobian(
+                state_trajectory[idx_u], u
+            )
+
+            state_trajectory[idx_u + 1] = state_next
+            final_input_trajectory[idx_u] = final_input
+            state_jacobian_trajectory[idx_u] = state_jacobian
+            final_input_jacobian_trajectory[idx_u] = final_input_jacobian
+
+        iterate = ILQRIterate(
+            state_trajectory=state_trajectory,  # type: ignore
+            input_trajectory=final_input_trajectory,  # type: ignore
+            state_jacobian_trajectory=state_jacobian_trajectory,  # type: ignore
+            input_jacobian_trajectory=final_input_jacobian_trajectory,  # type: ignore
+        )
+
+        return iterate
+
+    def _dynamics_and_jacobian(
+        self, current_state: DoubleMatrix, current_input: DoubleMatrix
+    ) -> Tuple[DoubleMatrix, DoubleMatrix, DoubleMatrix, DoubleMatrix]:
+        """
+        Propagates the state forward by one step and computes the corresponding state and input Jacobian matrices.
+        We also impose all constraints here to ensure the current input and next state are always feasible.
+        :param current_state: The current state z_k.
+        :param current_input: The applied input u_k.
+        :return: The next state z_{k+1}, (possibly modified) input u_k, and state (df/dz) and input (df/du) Jacobians.
+        """
+        x, y, heading, velocity, steering_angle = current_state
+
+        # Check steering angle is in expected range for valid Jacobian matrices.
+        assert (
+            np.abs(steering_angle) < np.pi / 2.0
+        ), f"The steering angle {steering_angle} is outside expected limits.  There is a singularity at delta = np.pi/2."
+
+        # Input constraints: clip inputs within bounds and then use.
+        current_input = self._clip_inputs(current_input)
+        acceleration, steering_rate = current_input
+
+        # Euler integration of bicycle model.
+        discretization_time = self._solver_params.discretization_time
+        wheelbase = self._solver_params.wheelbase
+
+        next_state: DoubleMatrix = np.copy(current_state)
+        next_state[0] += velocity * np.cos(heading) * discretization_time
+        next_state[1] += velocity * np.sin(heading) * discretization_time
+        next_state[2] += velocity * np.tan(steering_angle) / wheelbase * discretization_time
+        next_state[3] += acceleration * discretization_time
+        next_state[4] += steering_rate * discretization_time
+
+        # Constrain heading angle to lie within +/- pi.
+        next_state[2] = principal_value(next_state[2])
+
+        # State constraints: clip the steering_angle within bounds and update steering_rate accordingly.
+        next_steering_angle = self._clip_steering_angle(next_state[4])
+        applied_steering_rate = (next_steering_angle - steering_angle) / discretization_time
+        next_state[4] = next_steering_angle
+        current_input[1] = applied_steering_rate
+
+        # Now we construct and populate the state and input Jacobians.
+        state_jacobian: DoubleMatrix = np.eye(self._n_states, dtype=np.float64)
+        input_jacobian: DoubleMatrix = np.zeros((self._n_states, self._n_inputs), dtype=np.float64)
+
+        # Set a nonzero velocity to handle issues when linearizing at (near) zero velocity.
+        # This helps e.g. when the vehicle is stopped with zero steering angle and needs to accelerate/turn.
+        # Without this, the A matrix will indicate steering has no impact on heading due to Euler discretization.
+        # There will be a rank drop in the controllability matrix, so the discrete-time algebraic Riccati equation
+        # may not have a solution (uncontrollable subspace) or it may not be unique.
+        min_velocity_linearization = self._solver_params.min_velocity_linearization
+        if -min_velocity_linearization <= velocity and velocity <= min_velocity_linearization:
+            sign_velocity = 1.0 if velocity >= 0.0 else -1.0
+            velocity = sign_velocity * min_velocity_linearization
+
+        state_jacobian[0, 2] = -velocity * np.sin(heading) * discretization_time
+        state_jacobian[0, 3] = np.cos(heading) * discretization_time
+
+        state_jacobian[1, 2] = velocity * np.cos(heading) * discretization_time
+        state_jacobian[1, 3] = np.sin(heading) * discretization_time
+
+        state_jacobian[2, 3] = np.tan(steering_angle) / wheelbase * discretization_time
+        state_jacobian[2, 4] = velocity * discretization_time / (wheelbase * np.cos(steering_angle) ** 2)
+
+        input_jacobian[3, 0] = discretization_time
+        input_jacobian[4, 1] = discretization_time
+
+        return next_state, current_input, state_jacobian, input_jacobian
+
+    ####################################################################################################################
+    # Core LQR implementation.
+    ####################################################################################################################
+
+    def _run_lqr_backward_recursion(
+        self,
+        current_iterate: ILQRIterate,
+        reference_trajectory: DoubleMatrix,
+    ) -> ILQRInputPolicy:
+        """
+        Computes the locally optimal affine state feedback policy by applying dynamic programming to linear perturbation
+        dynamics about a specified linearization trajectory.  We include a trust region penalty as part of the cost.
+        :param current_iterate: Contains all relevant linearization information needed to compute LQR policy.
+        :param reference_trajectory: The desired state trajectory we are tracking.
+        :return: An affine state feedback policy - state feedback matrices and feedforward inputs found using LQR.
+        """
+        state_trajectory = current_iterate.state_trajectory
+        input_trajectory = current_iterate.input_trajectory
+        state_jacobian_trajectory = current_iterate.state_jacobian_trajectory
+        input_jacobian_trajectory = current_iterate.input_jacobian_trajectory
+
+        # Check reference matches the expected shape.
+        assert reference_trajectory.shape == state_trajectory.shape, "The reference trajectory has incorrect shape."
+
+        # Compute nominal error trajectory.
+        error_state_trajectory = state_trajectory - reference_trajectory
+        error_state_trajectory[:, 2] = principal_value(error_state_trajectory[:, 2])
+
+        # The value function has the form V_k(\Delta z_k) = \Delta z_k^T P_k \Delta z_k + 2 \rho_k^T \Delta z_k.
+        # So p_current = P_k is related to the Hessian of the value function at the current timestep.
+        # And rho_current = rho_k is part of the linear cost term in the value function at the current timestep.
+        p_current = self._state_cost_matrix + self._state_trust_region_cost_matrix
+        rho_current = self._state_cost_matrix @ error_state_trajectory[-1]
+
+        # The optimal LQR policy has the form \Delta u_k^* = K_k \Delta z_k + \kappa_k
+        # We refer to K_k as state_feedback_matrix and \kappa_k as feedforward input in the code below.
+        N = len(input_trajectory)
+        state_feedback_matrices = np.nan * np.ones((N, self._n_inputs, self._n_states), dtype=np.float64)
+        feedforward_inputs = np.nan * np.ones((N, self._n_inputs), dtype=np.float64)
+
+        for i in reversed(range(N)):
+            A = state_jacobian_trajectory[i]
+            B = input_jacobian_trajectory[i]
+            u = input_trajectory[i]
+            error = error_state_trajectory[i]
+
+            # Compute the optimal input policy for this timestep.
+            inverse_matrix_term = np.linalg.inv(
+                self._input_cost_matrix + self._input_trust_region_cost_matrix + B.T @ p_current @ B
+            )  # invertible since we checked input_cost / input_trust_region_cost are positive definite during creation.
+            state_feedback_matrix = -inverse_matrix_term @ B.T @ p_current @ A
+            feedforward_input = -inverse_matrix_term @ (self._input_cost_matrix @ u + B.T @ rho_current)
+
+            # Compute the optimal value function for this timestep.
+            a_closed_loop = A + B @ state_feedback_matrix
+
+            p_prior = (
+                self._state_cost_matrix
+                + self._state_trust_region_cost_matrix
+                + state_feedback_matrix.T @ self._input_cost_matrix @ state_feedback_matrix
+                + state_feedback_matrix.T @ self._input_trust_region_cost_matrix @ state_feedback_matrix
+                + a_closed_loop.T @ p_current @ a_closed_loop
+            )
+
+            rho_prior = (
+                self._state_cost_matrix @ error
+                + state_feedback_matrix.T @ self._input_cost_matrix @ (feedforward_input + u)
+                + state_feedback_matrix.T @ self._input_trust_region_cost_matrix @ feedforward_input
+                + a_closed_loop.T @ p_current @ B @ feedforward_input
+                + a_closed_loop.T @ rho_current
+            )
+
+            p_current = p_prior
+            rho_current = rho_prior
+
+            state_feedback_matrices[i] = state_feedback_matrix
+            feedforward_inputs[i] = feedforward_input
+
+        lqr_input_policy = ILQRInputPolicy(
+            state_feedback_matrices=state_feedback_matrices,  # type: ignore
+            feedforward_inputs=feedforward_inputs,  # type: ignore
+        )
+
+        return lqr_input_policy
+
+    def _update_inputs_with_policy(
+        self,
+        current_iterate: ILQRIterate,
+        lqr_input_policy: ILQRInputPolicy,
+    ) -> DoubleMatrix:
+        """
+        Used to update an iterate of iLQR by applying a perturbation input policy for local cost improvement.
+        :param current_iterate: Contains the state and input trajectory about which we linearized.
+        :param lqr_input_policy: Contains the LQR policy to apply.
+        :return: The next input trajectory found by applying the LQR policy.
+        """
+        state_trajectory = current_iterate.state_trajectory
+        input_trajectory = current_iterate.input_trajectory
+
+        # Trajectory of state perturbations while applying feedback policy.
+        # Starts with zero as the initial states match exactly, only later states might vary.
+        delta_state_trajectory = np.nan * np.ones((len(input_trajectory) + 1, self._n_states), dtype=np.float64)
+        delta_state_trajectory[0] = [0.0] * self._n_states
+
+        # This is the updated input trajectory we will return after applying the input perturbations.
+        input_next_trajectory = np.nan * np.ones_like(input_trajectory, dtype=np.float64)
+
+        zip_object = zip(
+            input_trajectory,
+            state_trajectory[:-1],
+            state_trajectory[1:],
+            lqr_input_policy.state_feedback_matrices,
+            lqr_input_policy.feedforward_inputs,
+        )
+
+        for input_idx, (input_lin, state_lin, state_lin_next, state_feedback_matrix, feedforward_input) in enumerate(
+            zip_object
+        ):
+            # Compute locally optimal input perturbation.
+            delta_state = delta_state_trajectory[input_idx]
+            delta_input = state_feedback_matrix @ delta_state + feedforward_input
+
+            # Apply state and input perturbation.
+            input_perturbed = input_lin + delta_input
+            state_perturbed = state_lin + delta_state
+            state_perturbed[2] = principal_value(state_perturbed[2])
+
+            # Run dynamics with perturbed state/inputs to get next state.
+            # We get the actually applied input since it might have been clipped/modified to satisfy constraints.
+            state_perturbed_next, input_perturbed, _, _ = self._dynamics_and_jacobian(state_perturbed, input_perturbed)
+
+            # Compute next state perturbation given next state.
+            delta_state_next = state_perturbed_next - state_lin_next
+            delta_state_next[2] = principal_value(delta_state_next[2])
+
+            delta_state_trajectory[input_idx + 1] = delta_state_next
+            input_next_trajectory[input_idx] = input_perturbed
+
+        assert ~np.any(np.isnan(input_next_trajectory)), "All next inputs should be valid float values."
+
+        return input_next_trajectory  # type: ignore

code/sim/ilqr/utils.py

@@ -0,0 +1,346 @@
+from typing import Tuple
+import numpy as np
+import numpy.typing as npt
+
+DoubleMatrix = npt.NDArray[np.float64]
+
+def principal_value(angle, min_=-np.pi):
+    """
+    Wrap heading angle in to specified domain (multiples of 2 pi alias),
+    ensuring that the angle is between min_ and min_ + 2 pi. This function raises an error if the angle is infinite
+    :param angle: rad
+    :param min_: minimum domain for angle (rad)
+    :return angle wrapped to [min_, min_ + 2 pi).
+    """
+    assert np.all(np.isfinite(angle)), "angle is not finite"
+    lhs = (angle - min_) % (2 * np.pi) + min_
+    return lhs
+
+def compute_steering_angle_feedback(
+    pose_reference, pose_current, lookahead_distance, k_lateral_error
+):
+    """
+    Given pose information, determines the steering angle feedback value to address initial tracking error.
+    This is based on the feedback controller developed in Section 2.2 of the following paper:
+    https://ddl.stanford.edu/publications/design-feedback-feedforward-steering-controller-accurate-path-tracking-and-stability
+    :param pose_reference: <np.ndarray: 3,> Contains the reference pose at the current timestep.
+    :param pose_current: <np.ndarray: 3,> Contains the actual pose at the current timestep.
+    :param lookahead_distance: [m] Distance ahead for which we should estimate lateral error based on a linear fit.
+    :param k_lateral_error: Feedback gain for lateral error used to determine steering angle feedback.
+    :return: [rad] The steering angle feedback to apply.
+    """
+    assert pose_reference.shape == (3,), "We expect a single reference pose."
+    assert pose_current.shape == (3,), "We expect a single current pose."
+
+    assert lookahead_distance > 0.0, "Lookahead distance should be positive."
+    assert k_lateral_error > 0.0, "Feedback gain for lateral error should be positive."
+
+    x_reference, y_reference, heading_reference = pose_reference
+    x_current, y_current, heading_current = pose_current
+
+    x_error = x_current - x_reference
+    y_error = y_current - y_reference
+    heading_error = principal_value(heading_current - heading_reference)
+
+    lateral_error = -x_error * np.sin(heading_reference) + y_error * np.cos(heading_reference)
+
+    return float(-k_lateral_error * (lateral_error + lookahead_distance * heading_error))
+
+def _convert_curvature_profile_to_steering_profile(
+    curvature_profile: DoubleMatrix,
+    discretization_time: float,
+    wheel_base: float,
+) -> Tuple[DoubleMatrix, DoubleMatrix]:
+    """
+    Converts from a curvature profile to the corresponding steering profile.
+    We assume a kinematic bicycle model where curvature = tan(steering_angle) / wheel_base.
+    For simplicity, we just use finite differences to determine steering rate.
+    :param curvature_profile: [rad] Curvature trajectory to convert.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param wheel_base: [m] The vehicle wheelbase parameter required for conversion.
+    :return: The [rad] steering angle and [rad/s] steering rate (derivative) profiles.
+    """
+    assert discretization_time > 0.0, "Discretization time must be positive."
+    assert wheel_base > 0.0, "The vehicle's wheelbase length must be positive."
+
+    steering_angle_profile = np.arctan(wheel_base * curvature_profile)
+    steering_rate_profile = np.diff(steering_angle_profile) / discretization_time
+
+    return steering_angle_profile, steering_rate_profile
+
+
+def _get_xy_heading_displacements_from_poses(poses: DoubleMatrix) -> Tuple[DoubleMatrix, DoubleMatrix]:
+    """
+    Returns position and heading displacements given a pose trajectory.
+    :param poses: <np.ndarray: num_poses, 3> A trajectory of poses (x, y, heading).
+    :return: Tuple of xy displacements with shape (num_poses-1, 2) and heading displacements with shape (num_poses-1,).
+    """
+    assert len(poses.shape) == 2, "Expect a 2D matrix representing a trajectory of poses."
+    assert poses.shape[0] > 1, "Cannot get displacements given an empty or single element pose trajectory."
+    assert poses.shape[1] == 3, "Expect pose to have three elements (x, y, heading)."
+
+    # Compute displacements that are used to complete the kinematic state and input.
+    pose_differences = np.diff(poses, axis=0)
+    xy_displacements = pose_differences[:, :2]
+    heading_displacements = principal_value(pose_differences[:, 2])
+
+    return xy_displacements, heading_displacements
+
+
+def _make_banded_difference_matrix(number_rows: int) -> DoubleMatrix:
+    """
+    Returns a banded difference matrix with specified number_rows.
+    When applied to a vector [x_1, ..., x_N], it returns [x_2 - x_1, ..., x_N - x_{N-1}].
+    :param number_rows: The row dimension of the banded difference matrix (e.g. N-1 in the example above).
+    :return: A banded difference matrix with shape (number_rows, number_rows+1).
+    """
+    banded_matrix: DoubleMatrix = -1.0 * np.eye(number_rows + 1, dtype=np.float64)[:-1, :]
+    for ind in range(len(banded_matrix)):
+        banded_matrix[ind, ind + 1] = 1.0
+
+    return banded_matrix
+
+
+
+def _fit_initial_velocity_and_acceleration_profile(
+    xy_displacements: DoubleMatrix, heading_profile: DoubleMatrix, discretization_time: float, jerk_penalty: float
+) -> Tuple[float, DoubleMatrix]:
+    """
+    Estimates initial velocity (v_0) and acceleration ({a_0, ...}) using least squares with jerk penalty regularization.
+    :param xy_displacements: [m] Deviations in x and y occurring between M+1 poses, a M by 2 matrix.
+    :param heading_profile: [rad] Headings associated to the starting timestamp for xy_displacements, a M-length vector.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param jerk_penalty: A regularization parameter used to penalize acceleration differences.  Should be positive.
+    :return: Least squares solution for initial velocity (v_0) and acceleration profile ({a_0, ..., a_M-1})
+             for M displacement values.
+    """
+    assert discretization_time > 0.0, "Discretization time must be positive."
+    assert jerk_penalty > 0, "Should have a positive jerk_penalty."
+
+    assert len(xy_displacements.shape) == 2, "Expect xy_displacements to be a matrix."
+    assert xy_displacements.shape[1] == 2, "Expect xy_displacements to have 2 columns."
+
+    num_displacements = len(xy_displacements)  # aka M in the docstring
+
+    assert heading_profile.shape == (
+        num_displacements,
+    ), "Expect the length of heading_profile to match that of xy_displacements."
+
+    # Core problem: minimize_x ||y-Ax||_2
+    y = xy_displacements.flatten()  # Flatten to a vector, [delta x_0, delta y_0, ...]
+
+    A: DoubleMatrix = np.zeros((2 * num_displacements, num_displacements), dtype=np.float64)
+    for idx_timestep, heading in enumerate(heading_profile):
+        start_row = 2 * idx_timestep  # Which row of A corresponds to x-coordinate information at timestep k.
+
+        # Related to v_0, initial velocity - column 0.
+        # We fill in rows for measurements delta x_k, delta y_k.
+        A[start_row : (start_row + 2), 0] = np.array(
+            [
+                np.cos(heading) * discretization_time,
+                np.sin(heading) * discretization_time,
+            ],
+            dtype=np.float64,
+        )
+
+        if idx_timestep > 0:
+            # Related to {a_0, ..., a_k-1}, acceleration profile - column 1 to k.
+            # We fill in rows for measurements delta x_k, delta y_k.
+            A[start_row : (start_row + 2), 1 : (1 + idx_timestep)] = np.array(
+                [
+                    [np.cos(heading) * discretization_time**2],
+                    [np.sin(heading) * discretization_time**2],
+                ],
+                dtype=np.float64,
+            )
+
+    # Regularization using jerk penalty, i.e. difference of acceleration values.
+    # If there are M displacements, then we have M - 1 acceleration values.
+    # That means we have M - 2 jerk values, thus we make a banded difference matrix of that size.
+    banded_matrix = _make_banded_difference_matrix(num_displacements - 2)
+    R: DoubleMatrix = np.block([np.zeros((len(banded_matrix), 1)), banded_matrix])
+
+    # Compute regularized least squares solution.
+    x = np.linalg.pinv(A.T @ A + jerk_penalty * R.T @ R) @ A.T @ y
+
+    # Extract profile from solution.
+    initial_velocity = x[0]
+    acceleration_profile = x[1:]
+
+    return initial_velocity, acceleration_profile
+
+
+def _generate_profile_from_initial_condition_and_derivatives(
+    initial_condition: float, derivatives: DoubleMatrix, discretization_time: float
+) -> DoubleMatrix:
+    """
+    Returns the corresponding profile (i.e. trajectory) given an initial condition and derivatives at
+    multiple timesteps by integration.
+    :param initial_condition: The value of the variable at the initial timestep.
+    :param derivatives: The trajectory of time derivatives of the variable at timesteps 0,..., N-1.
+    :param discretization_time: [s] Time discretization used for integration.
+    :return: The trajectory of the variable at timesteps 0,..., N.
+    """
+    assert discretization_time > 0.0, "Discretization time must be positive."
+
+    profile = initial_condition + np.insert(np.cumsum(derivatives * discretization_time), 0, 0.0)
+
+    return profile  # type: ignore
+
+
+def _fit_initial_curvature_and_curvature_rate_profile(
+    heading_displacements: DoubleMatrix,
+    velocity_profile: DoubleMatrix,
+    discretization_time: float,
+    curvature_rate_penalty: float,
+    initial_curvature_penalty: float = 1e-10,
+) -> Tuple[float, DoubleMatrix]:
+    """
+    Estimates initial curvature (curvature_0) and curvature rate ({curvature_rate_0, ...})
+    using least squares with curvature rate regularization.
+    :param heading_displacements: [rad] Angular deviations in heading occuring between timesteps.
+    :param velocity_profile: [m/s] Estimated or actual velocities at the timesteps matching displacements.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate.  Should be positive.
+    :param initial_curvature_penalty: A regularization parameter to handle zero initial speed.  Should be positive and small.
+    :return: Least squares solution for initial curvature (curvature_0) and curvature rate profile
+             (curvature_rate_0, ..., curvature_rate_{M-1}) for M heading displacement values.
+    """
+    assert discretization_time > 0.0, "Discretization time must be positive."
+    assert curvature_rate_penalty > 0.0, "Should have a positive curvature_rate_penalty."
+    assert initial_curvature_penalty > 0.0, "Should have a positive initial_curvature_penalty."
+
+    # Core problem: minimize_x ||y-Ax||_2
+    y = heading_displacements
+    A: DoubleMatrix = np.tri(len(y), dtype=np.float64)  # lower triangular matrix
+    A[:, 0] = velocity_profile * discretization_time
+
+    for idx, velocity in enumerate(velocity_profile):
+        if idx == 0:
+            continue
+        A[idx, 1:] *= velocity * discretization_time**2
+
+    # Regularization on curvature rate.  We add a small but nonzero weight on initial curvature too.
+    # This is since the corresponding row of the A matrix might be zero if initial speed is 0, leading to singularity.
+    # We guarantee that Q is positive definite such that the minimizer of the least squares problem is unique.
+    Q: DoubleMatrix = curvature_rate_penalty * np.eye(len(y))
+    Q[0, 0] = initial_curvature_penalty
+
+    # Compute regularized least squares solution.
+    x = np.linalg.pinv(A.T @ A + Q) @ A.T @ y
+
+    # Extract profile from solution.
+    initial_curvature = x[0]
+    curvature_rate_profile = x[1:]
+
+    return initial_curvature, curvature_rate_profile
+
+
+def get_velocity_curvature_profiles_with_derivatives_from_poses(
+    discretization_time: float,
+    poses: DoubleMatrix,
+    jerk_penalty: float,
+    curvature_rate_penalty: float,
+) -> Tuple[DoubleMatrix, DoubleMatrix, DoubleMatrix, DoubleMatrix]:
+    """
+    Main function for joint estimation of velocity, acceleration, curvature, and curvature rate given N poses
+    sampled at discretization_time.  This is done by solving two least squares problems with the given penalty weights.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param poses: <np.ndarray: num_poses, 3> A trajectory of N poses (x, y, heading).
+    :param jerk_penalty: A regularization parameter used to penalize acceleration differences.  Should be positive.
+    :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate.  Should be positive.
+    :return: Profiles for velocity (N-1), acceleration (N-2), curvature (N-1), and curvature rate (N-2).
+    """
+    xy_displacements, heading_displacements = _get_xy_heading_displacements_from_poses(poses)
+
+    # Compute initial velocity + acceleration least squares solution and extract results.
+    # Note: If we have M displacements, we require the M associated heading values.
+    #       Therefore, we exclude the last heading in the call below.
+    initial_velocity, acceleration_profile = _fit_initial_velocity_and_acceleration_profile(
+        xy_displacements=xy_displacements,
+        heading_profile=poses[:-1, 2],
+        discretization_time=discretization_time,
+        jerk_penalty=jerk_penalty,
+    )
+
+    velocity_profile = _generate_profile_from_initial_condition_and_derivatives(
+        initial_condition=initial_velocity,
+        derivatives=acceleration_profile,
+        discretization_time=discretization_time,
+    )
+
+    # Compute initial curvature + curvature rate least squares solution and extract results.  It relies on velocity fit.
+    initial_curvature, curvature_rate_profile = _fit_initial_curvature_and_curvature_rate_profile(
+        heading_displacements=heading_displacements,
+        velocity_profile=velocity_profile,
+        discretization_time=discretization_time,
+        curvature_rate_penalty=curvature_rate_penalty,
+    )
+
+    curvature_profile = _generate_profile_from_initial_condition_and_derivatives(
+        initial_condition=initial_curvature,
+        derivatives=curvature_rate_profile,
+        discretization_time=discretization_time,
+    )
+
+    return velocity_profile, acceleration_profile, curvature_profile, curvature_rate_profile
+
+
+
+def complete_kinematic_state_and_inputs_from_poses(
+    discretization_time: float,
+    wheel_base: float,
+    poses: DoubleMatrix,
+    jerk_penalty: float,
+    curvature_rate_penalty: float,
+) -> Tuple[DoubleMatrix, DoubleMatrix]:
+    """
+    Main function for joint estimation of velocity, acceleration, steering angle, and steering rate given poses
+    sampled at discretization_time and the vehicle wheelbase parameter for curvature -> steering angle conversion.
+    One caveat is that we can only determine the first N-1 kinematic states and N-2 kinematic inputs given
+    N-1 displacement/difference values, so we need to extrapolate to match the length of poses provided.
+    This is handled by repeating the last input and extrapolating the motion model for the last state.
+    :param discretization_time: [s] Time discretization used for integration.
+    :param wheel_base: [m] The wheelbase length for the kinematic bicycle model being used.
+    :param poses: <np.ndarray: num_poses, 3> A trajectory of poses (x, y, heading).
+    :param jerk_penalty: A regularization parameter used to penalize acceleration differences.  Should be positive.
+    :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate.  Should be positive.
+    :return: kinematic_states (x, y, heading, velocity, steering_angle) and corresponding
+            kinematic_inputs (acceleration, steering_rate).
+    """
+    (
+        velocity_profile,
+        acceleration_profile,
+        curvature_profile,
+        curvature_rate_profile,
+    ) = get_velocity_curvature_profiles_with_derivatives_from_poses(
+        discretization_time=discretization_time,
+        poses=poses,
+        jerk_penalty=jerk_penalty,
+        curvature_rate_penalty=curvature_rate_penalty,
+    )
+
+    # Convert to steering angle given the wheelbase parameter.  At this point, we don't need to worry about curvature.
+    steering_angle_profile, steering_rate_profile = _convert_curvature_profile_to_steering_profile(
+        curvature_profile=curvature_profile,
+        discretization_time=discretization_time,
+        wheel_base=wheel_base,
+    )
+
+    # Extend input fits with a repeated element and extrapolate state fits to match length of poses.
+    # This is since we fit with N-1 displacements but still have N poses at the end to deal with.
+    acceleration_profile = np.append(acceleration_profile, acceleration_profile[-1])
+    steering_rate_profile = np.append(steering_rate_profile, steering_rate_profile[-1])
+
+    velocity_profile = np.append(
+        velocity_profile, velocity_profile[-1] + acceleration_profile[-1] * discretization_time
+    )
+    steering_angle_profile = np.append(
+        steering_angle_profile, steering_angle_profile[-1] + steering_rate_profile[-1] * discretization_time
+    )
+
+    # Collect completed state and input in matrices.
+    kinematic_states: DoubleMatrix = np.column_stack((poses, velocity_profile, steering_angle_profile))
+    kinematic_inputs: DoubleMatrix = np.column_stack((acceleration_profile, steering_rate_profile))
+
+    return kinematic_states, kinematic_inputs

code/sim/pyproject.toml

@@ -0,0 +1,9 @@
+[tool.hatch.build.targets.wheel]
+packages = ["hugsim-env"]
+
+[project]
+name = "hugsim-env"
+version = "0.0.1"
+dependencies = [
+  "gymnasium",
+]

code/sim/setup.py

@@ -0,0 +1,7 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="hugsim-env",
+    version="0.0.1",
+    packages=find_packages(),
+)

code/sim/utils/agent_controller.py

@@ -0,0 +1,323 @@
+import math
+import random
+import numpy as np
+from trajdata.maps import VectorMap
+from submodules.Pplan.Sampling.spline_planner import SplinePlanner
+import torch
+import time
+import math
+from copy import deepcopy
+from utils.dynamic_utils import unicycle
+
+
+def constant_tracking(state, path, dt):
+    '''
+    Args:
+        state: current state of the vehicle, of size [x, y, yaw, speed]
+        path: the path to follow, of size (N, [x, y, yaw])
+        dt: time duration
+    '''
+
+    # find the nearest point in the path
+    dists = torch.norm(path[:, :2] - state[None, :2], dim=1)
+    nearest_index = torch.argmin(dists)
+
+    # find the target point
+    lookahead_distance = state[3] * dt
+    target = path[-1]
+    is_end = True
+    for i in range(nearest_index + 1, len(path)):
+        if torch.norm(path[i, :2] - state[:2]) > lookahead_distance:
+            target = path[i]
+            is_end = False
+            break
+
+    # compute the new state
+    target_distance = torch.norm(target[:2] - state[:2])
+    ratio = lookahead_distance / target_distance.clamp(min=1e-6)
+    ratio = ratio.clamp(max=1.0)
+
+    new_state = deepcopy(state)
+    new_state[:2] = state[:2] + ratio * (target[:2] - state[:2])
+    new_state[2] = torch.atan2(
+        state[2].sin() + ratio * (target[2].sin() - state[2].sin()),
+        state[2].cos() + ratio * (target[2].cos() - state[2].cos())
+    )
+    if is_end:
+        new_state[3] = 0
+
+    return new_state
+
+
+def constant_headaway(states, num_steps, dt):
+    '''
+    Args:
+        states: current states of a batch of vehicles, of size (num_agents, [x, y, yaw, speed])
+        num_steps: number of steps to move forward
+        dt: time duration
+    Return:
+        trajs: the trajectories of the vehicles, of size (num_agents, num_steps, [x, y, yaw, speed])
+    '''
+
+    # state: [x, y, yaw, speed]
+    x = states[:, 0]
+    y = states[:, 1]
+    yaw = states[:, 2]
+    speed = states[:, 3]
+
+    # Generate time steps
+    t_steps = torch.arange(num_steps) * dt
+
+    # Calculate dx and dy for each step
+    dx = torch.outer(speed * torch.sin(yaw), t_steps)
+    dy = torch.outer(speed * torch.cos(yaw), t_steps)
+
+    # Update x and y positions
+    x_traj = x.unsqueeze(1) + dx
+    y_traj = y.unsqueeze(1) + dy
+
+    # Replicate the yaw and speed for each time step
+    yaw_traj = yaw.unsqueeze(1).repeat(1, num_steps)
+    speed_traj = speed.unsqueeze(1).repeat(1, num_steps)
+
+    # Stack the x, y, yaw, and speed components to form the trajectory
+    trajs = torch.stack((x_traj, y_traj, yaw_traj, speed_traj), dim=-1)
+
+    return trajs
+
+
+class IDM:
+    def __init__(
+            self, v0=30.0, s0=5.0, T=2.0, a=2.0, b=4.0, delta=4.0,
+            lookahead_path_length=100, lead_distance_threshold=1.0
+    ):
+        '''
+        Args:
+            v0: desired speed
+            s0: minimum gap
+            T: safe time headway
+            a: max acceleration
+            b: comfortable deceleration
+            delta: acceleration exponent
+            lookahead_path_length: the length of path to look ahead
+            lead_distance_threshold: the distance to consider a vehicle as a lead vehicle
+        '''
+        self.v0 = v0
+        self.s0 = s0
+        self.T = T
+        self.a = a
+        self.b = b
+        self.delta = delta
+        self.lookahead_path_length = lookahead_path_length
+        self.lead_distance_threshold = lead_distance_threshold
+
+    def update(self, state, path, dt, neighbors):
+        '''
+        Args:
+            state: current state of the vehicle, of size [x, y, yaw, speed]
+            path: the path to follow, of size (N, [x, y, yaw])
+            dt: time duration
+            neighbors: the future states of the neighbors, of size (K, T, [x, y, yaw, speed])
+        '''
+
+        if path is None:
+            return deepcopy(state)
+
+        # find the nearest point in the path
+        dists = torch.norm(path[:, :2] - state[None, :2], dim=1)
+        nearest_index = torch.argmin(dists)
+
+
+        # lookahead_distance = state[3] * dt
+        # lookahead_targe = state[:2] + np.array([np.sin(state[2]) * lookahead_distance, np.cos(state[2]) * lookahead_distance])
+        # # target = path[-1]
+        # is_end = False
+        # target_idx = torch.argmin(torch.norm(path[:, :2] - lookahead_targe, dim=-1))
+        # target = path[target_idx]
+
+        # find the target point
+        lookahead_distance = state[3] * dt
+        target = path[-1]
+        is_end = True
+        for i in range(nearest_index + 1, len(path)):
+            if torch.norm(path[i, :2] - state[:2]) > lookahead_distance:
+                target = path[i]
+                is_end = False
+                break
+
+        # distance between neighbors and the path
+        lookahead_path = path[nearest_index + 1:][:self.lookahead_path_length]
+        lookahead_neighbors = neighbors[..., None, :].expand(
+            -1, -1, lookahead_path.shape[0], -1
+        )  # (K, T, n, 4)
+
+        dists_neighbors = torch.norm(
+            lookahead_neighbors[..., :2] - lookahead_path[None, None, :, :2], dim=-1
+        )  # (K, T, n)
+        indices_neighbors = torch.arange(
+            lookahead_path.shape[0]
+        )[None, None].expand_as(dists_neighbors)
+
+        # determine lead vehicles
+        is_lead = (dists_neighbors < self.lead_distance_threshold)
+        if is_lead.any():
+            # compute lead distance
+            indices_lead = indices_neighbors[is_lead]  # (num_lead)
+            lookahead_lengths = torch.cumsum(torch.norm(
+                lookahead_path[1:, :2] - lookahead_path[:-1, :2], dim=1
+            ), dim=0)
+            lookahead_lengths = torch.cat([lookahead_lengths, lookahead_lengths[-1:]])
+            lead_distance = lookahead_lengths[indices_lead]
+
+            # compute lead speed
+            states_lead = lookahead_neighbors[is_lead]  # (num_lead, 4)
+            ori_speed_lead = states_lead[:, 3]
+            yaw_lead = states_lead[:, 2]
+            yaw_path = lookahead_path[indices_lead, 2]
+            lead_speed = ori_speed_lead * (yaw_lead - yaw_path).cos()
+
+            # compute acceleration
+            ego_speed = state[3]
+            delta_v = ego_speed - lead_speed
+            s_star = self.s0 + \
+                     (ego_speed * self.T + ego_speed * delta_v / (2 * math.sqrt(self.a * self.b))).clamp(min=0)
+            acceleration = self.a * (1 - (ego_speed / self.v0) ** self.delta - (s_star / lead_distance) ** 2)
+            acceleration = acceleration.min()
+        else:
+            acceleration = self.a * (1 - (state[3] / self.v0) ** self.delta)
+
+        # compute the new state
+        target_distance = torch.norm(target[:2] - state[:2])
+        ratio = lookahead_distance / target_distance.clamp(min=1e-6)
+        ratio = ratio.clamp(max=1.0)
+
+        new_state = deepcopy(state)
+        new_state[:2] = state[:2] + ratio * (target[:2] - state[:2])
+        new_state[2] = torch.atan2(
+            state[2].sin() + ratio * (target[2].sin() - state[2].sin()),
+            state[2].cos() + ratio * (target[2].cos() - state[2].cos())
+        )
+        if is_end:
+            new_state[3] = 0
+        else:
+            new_state[3] = (state[3] + acceleration * dt).clamp(min=0)
+
+        return new_state
+
+
+class AttackPlanner:
+    def __init__(self, pred_steps=20, ATTACK_FREQ = 3, best_k=1, device='cpu'):
+        self.device = device
+        self.predict_steps = pred_steps
+        self.best_k = best_k
+
+        self.planner = SplinePlanner(
+            device,
+            N_seg=self.predict_steps,
+            acce_grid=torch.linspace(-2, 5, 10).to(self.device),
+            acce_bound=[-6, 5],
+            vbound=[-2, 50]
+        )
+        self.planner.psi_bound = [-math.pi * 2, math.pi * 2]
+
+        self.exec_traj = None
+        self.exec_pointer = 1
+
+    def update(
+            self, state, unified_map, dt,
+            neighbors, attacked_states,
+            new_plan=True
+    ):
+        '''
+        Args:
+            state: current state of the vehicle, of size [x, y, yaw, speed]
+            vector_map: the vector map
+            attacked_states: future states of the attacked agent, of size (T, [x, y, yaw, speed])
+            neighbors: future states of the neighbors, of size (K, T, [x, y, yaw, speed])
+            new_plan: whether to generate a new plan
+        '''
+        assert self.exec_pointer > 0
+
+        # directly execute the current plan
+        if not new_plan:
+            if self.exec_traj is not None and \
+                    self.exec_pointer < self.exec_traj.shape[0]:
+                next_state = self.exec_traj[self.exec_pointer]
+                self.exec_pointer += 1
+                return next_state
+            else:
+                new_plan = True
+
+        assert attacked_states.shape[0] == self.predict_steps
+
+        # state: [x, y, yaw, speed]
+        x, y, yaw, speed = state
+
+        # query vector map to get lanes
+        query_xyzr = np.array([x, y, 0, yaw + np.pi / 2])
+        # query_xyzr = unified_map.xyzr_local2world(np.array([x, y, 0, yaw]))
+        # lanes = unified_map.vector_map.get_lanes_within(query_xyzr[:3], dist=30)
+        # lanes = [unified_map.batch_xyzr_world2local(l.center.xyzh)[:, [0,1,3]] for l in lanes]
+        # lanes = [l.center.xyzh[:, [0,1,3]] for l in lanes]
+        lanes = None
+
+        # for lane in lanes:
+        #     plt.plot(lane[:, 0], lane[:, 1], 'k--', linewidth=0.5, alpha=0.5)
+
+        # generate spline trajectories
+        x0 = torch.tensor([query_xyzr[0], query_xyzr[1], speed, query_xyzr[3]], device=self.device)
+        possible_trajs, xf_set = self.planner.gen_trajectories(x0, self.predict_steps * dt, lanes,
+                                                               dyn_filter=True)  # (num_trajs, T-1, [x, y, v, a, yaw, r, t])
+        if possible_trajs.shape[0] == 0:
+            trajs = constant_headaway(state[None], self.predict_steps, dt)  # (1, T, [x, y, yaw, speed])
+        else:
+            trajs = torch.cat([
+                state[None, None].expand(possible_trajs.shape[0], -1, -1),
+                possible_trajs[..., [0, 1, 4, 2]]
+            ], dim=1)
+
+        # select the best trajectory
+        attack_distance = torch.norm(attacked_states[None, :, :2] - trajs[..., :2], dim=-1)
+        cost_attack = attack_distance.min(dim=1).values
+        cost_collision = (
+                    torch.norm(neighbors[None, ..., :2] - trajs[:, None, :, :2], dim=-1).min(dim=-1).values < 2.0).sum(
+            dim=-1)
+        cost = cost_attack + 0.1 * cost_collision
+        values, indices = torch.topk(cost, self.best_k, largest=False)
+        random_index = torch.randint(0, self.best_k, (1,)).item()
+        selected_index = indices[random_index]
+        traj_best = trajs[selected_index]
+
+        # produce next state
+        self.exec_traj = traj_best
+        self.exec_traj[:, 2] -= np.pi / 2
+        self.exec_pointer = 1
+        next_state = self.exec_traj[self.exec_pointer]
+        # next_state[0] = -next_state[0]
+        self.exec_pointer += 1
+
+        return next_state
+
+
+class ConstantPlanner:
+    def __init__(self):
+        return
+
+    def update(self, state, dt):
+        a, b, yaw, v = state
+        a = a - v * np.sin(yaw) * dt
+        b = b + v * np.cos(yaw) * dt
+        return torch.tensor([a, b, yaw, v])
+    
+
+class UnicyclePlanner:
+    def __init__(self, uc_path, speed=1.0):
+        self.uc_model = unicycle.restore(torch.load(uc_path, weights_only=False))
+        self.t = 0
+        self.speed = speed
+    
+    def update(self, dt):
+        self.t += dt * self.speed
+        a, b, v, pitchroll, yaw, h = self.uc_model.forward(self.t)
+        # return torch.tensor([a, b, yaw, v]), pitchroll.detach().cpu(), h.item()
+        return torch.tensor([a, b, yaw, v])

code/sim/utils/launch_ad.py

@@ -0,0 +1,30 @@
+import subprocess
+import time
+import os
+
+
+def launch(shell_path, cuda_id, output):
+    os.makedirs(output, exist_ok=True)
+    print(os.path.join(output, 'output.txt'))
+    print(shell_path, cuda_id, output)
+    with open(os.path.join(output, 'output.txt'), 'w') as f:
+        process = subprocess.Popen(
+            ["zsh", shell_path, cuda_id, output], stdout=f, stderr=f
+        )
+    return process
+
+
+def check_alive(process, tolerant=100):
+    i = 0
+    while i < tolerant:
+        return_code = process.poll()
+        if return_code is not None:
+            print(f"The AD algorithm completed with return code {return_code}.")
+            process.kill()
+            return
+        elif i % 5 == 0:
+            print(f"The AD algorithm is still running, remaining tolerant {tolerant - i}.")
+        time.sleep(1)
+        i += 1
+    process.kill()
+    print("The AD algorithm process is killed.")

code/sim/utils/plan.py

@@ -0,0 +1,238 @@
+import numpy as np
+import torch
+from scipy.spatial.transform import Rotation as SCR
+import roma
+from collections import namedtuple
+from sim.utils.agent_controller import constant_headaway
+from sim.utils import agent_controller
+from collections import defaultdict
+from trajdata import AgentType, UnifiedDataset
+from trajdata.maps import MapAPI
+from trajdata.simulation import SimulationScene
+from sim.utils.sim_utils import rt2pose, pose2rt
+from sim.utils.agent_controller import IDM, AttackPlanner, ConstantPlanner, UnicyclePlanner
+import os
+import json
+
+Model = namedtuple('Models', ['model_path', 'controller', 'controller_args'])
+
+
+class planner:
+    def __init__(self, plan_list, scene_path=None, dt=0.2, unified_map=None, ground=None):
+        self.unified_map = unified_map
+        self.ground = ground
+        self.PREDICT_STEPS = 20
+        self.NUM_NEIGHBORS = 3
+        
+        self.rectify_angle = 0
+        if self.unified_map is not None:
+            self.rectify_angle = self.unified_map.rectify_angle
+        
+        # plan_list: a, b, height, yaw, v, model_path, controller, controller_args: dict
+        self.stats, self.route, self.controller, self.ckpts, self.wlhs = {}, {}, {}, {}, {}
+        self.dt = dt
+        self.ATTACK_FREQ = 3
+        for iid, args in enumerate(plan_list):
+            if args[6] == "UnicyclePlanner":
+                # self.ckpts[f"agent_{iid}"] = os.path.join(scene_path, "ckpts", f"dynamic_{args[7]}_chkpnt30000.pth")
+                # self.wlhs[f'agent_{iid}'] = [2.0, 4.0, 1.5]
+                self.ckpts[f'agent_{iid}'] = os.path.join(args[5], 'gs.pth')
+                with open(os.path.join(args[5], 'wlh.json')) as f:
+                    self.wlhs[f'agent_{iid}'] = json.load(f)
+                uc_configs = args[7]
+                self.controller[f"agent_{iid}"] = UnicyclePlanner(os.path.join(scene_path, f"unicycle_{uc_configs['uc_id']}.pth"), speed=uc_configs['speed'])
+                a, b, v, pitchroll, yaw, h = self.controller[f"agent_{iid}"].uc_model.forward(0.0)
+                self.stats[f'agent_{iid}'] = torch.tensor([a, b, args[2], yaw, v])
+                self.route[f'agent_{iid}'] = None
+            else:
+                model = Model(*args[5:])
+                self.stats[f'agent_{iid}'] = torch.tensor(args[:5])  # a, b, height, yaw, v
+                self.stats[f'agent_{iid}'][3] += self.rectify_angle
+                self.route[f'agent_{iid}'] = None
+                self.ckpts[f'agent_{iid}'] = os.path.join(model.model_path, 'gs.pth')
+                with open(os.path.join(model.model_path, 'wlh.json')) as f:
+                    self.wlhs[f'agent_{iid}'] = json.load(f)
+                self.controller[f'agent_{iid}'] = getattr(agent_controller, model.controller)(**model.controller_args)
+                if model.controller == "AttackPlanner":
+                    self.ATTACK_FREQ = model.controller_args["ATTACK_FREQ"]
+
+    def update_ground(self, ground):
+        self.ground = ground
+
+    def update_agent_route(self):
+        assert self.unified_map is not None, "Map shouldn't be None to forecast agent path"
+        for iid, stat in self.stats.items():
+            path = self.unified_map.get_route(stat)
+            if path is None:
+                print("path not found at ", self.stats)
+            if path is not None:
+                self.route[iid] = torch.from_numpy(np.hstack([path[:, :2], path[:, -1:]]))
+
+    def ground_height(self, u, v):
+        cam_poses, cam_height, _ = self.ground
+        cam_poses = torch.from_numpy(cam_poses)
+        cam_dist = np.sqrt(
+            (cam_poses[:-1, 0, 3] - u) ** 2 + (cam_poses[:-1, 2, 3] - v) ** 2
+        )
+        nearest_cam_idx = np.argmin(cam_dist, axis=0)
+        nearest_c2w = cam_poses[nearest_cam_idx]
+        
+        nearest_w2c = np.linalg.inv(nearest_c2w)
+        uv_local = nearest_w2c[:3, :3] @ np.array([u, 0, v]) + nearest_w2c[:3, 3]
+        uv_local[1] = 0
+        uv_world = nearest_c2w[:3, :3] @ uv_local + nearest_c2w[:3, 3]
+        
+        return uv_world[1] + cam_height
+
+    def plan_traj(self, t, ego_stats):
+        all_stats = [ego_stats]
+        for iid, stat in self.stats.items():
+            all_stats.append(stat[[0, 1, 3, 4]])  # a, b, yaw, v
+        all_stats = torch.stack(all_stats, dim=0)
+        future_states = constant_headaway(all_stats, num_steps=self.PREDICT_STEPS, dt=self.dt)
+
+        b2ws = {}
+        for iid, stat in self.stats.items():
+            # find closet neighbors
+            curr_xy_agents = all_stats[:, :2]
+            distance_agents = torch.norm(curr_xy_agents - stat[:2], dim=-1)
+            neighbor_idx = torch.argsort(distance_agents)[1:self.NUM_NEIGHBORS + 1]
+            neighbors = future_states[neighbor_idx]
+
+            controller = self.controller[iid]
+            if type(controller) is IDM:
+                next_xyrv = controller.update(state=stat[[0, 1, 3, 4]], path=self.route[iid], dt=self.dt,
+                                              neighbors=neighbors)
+            elif type(controller) is AttackPlanner:
+                safe_neighbors = neighbors[1:, ...]
+                next_xyrv = controller.update(state=stat[[0, 1, 3, 4]], unified_map=self.unified_map, dt=0.1,
+                                              neighbors=safe_neighbors, attacked_states=future_states[0],
+                                              new_plan=((t // self.dt) % self.ATTACK_FREQ == 0))
+            elif type(controller) is ConstantPlanner:
+                next_xyrv = controller.update(state=stat[[0, 1, 3, 4]], dt=self.dt)
+            elif type(controller) is UnicyclePlanner:
+                next_xyrv = controller.update(dt=self.dt)
+            else:
+                raise NotImplementedError
+            next_stat = torch.zeros_like(stat)
+            next_stat[[0, 1, 3, 4]] = next_xyrv.float()
+            next_stat[2] = stat[2]
+            self.stats[iid] = next_stat
+
+            b2w = np.eye(4)
+            h = self.ground_height(next_xyrv[0].numpy(), next_xyrv[1].numpy())
+            if type(controller) is UnicyclePlanner:
+                # b2w[:3, :3] = SCR.from_euler('xzy', [pitch_roll[0], pitch_roll[1], stat[3]]).as_matrix()
+                b2w[:3, :3] = SCR.from_euler('y', [-stat[3]]).as_matrix()
+                b2w[:3, 3] = np.array([next_stat[0], h + stat[2], next_stat[1]])
+            else:
+                b2w[:3, :3] = SCR.from_euler('y', [-stat[3] - np.pi / 2 - self.rectify_angle]).as_matrix()
+                b2w[:3, 3] = np.array([next_stat[0], h + stat[2], next_stat[1]])
+            b2ws[iid] = torch.tensor(b2w).float().cuda()
+
+        return [b2ws, {}]
+
+
+class UnifiedMap:
+    def __init__(self, datapath, version, scene_name):
+        self.datapath = datapath
+        self.version = version
+
+        self.dataset = UnifiedDataset(
+            desired_data=[self.version],
+            data_dirs={
+                self.version: self.datapath,
+            },
+            cache_location="/app/app_datas/nusc_map_cache",
+            only_types=[AgentType.VEHICLE],
+            agent_interaction_distances=defaultdict(lambda: 50.0),
+            desired_dt=0.1,
+            num_workers=4,
+            verbose=True,
+        )
+
+        self.map_api = MapAPI(self.dataset.cache_path)
+
+        self.scene = None
+        for scene in list(self.dataset.scenes()):
+            if scene.name == scene_name:
+                self.scene = scene
+        assert self.scene is not None, f"Can't find scene {scene_name}"
+        self.vector_map = self.map_api.get_map(
+            f"{self.version}:{self.scene.location}"
+        )
+        self.ego_start_pos, self.ego_start_yaw = self.get_start_pose()
+        self.rectify_angle = 0
+        if self.ego_start_yaw < 0:
+            self.ego_start_yaw += np.pi
+            self.rectify_angle = np.pi
+        self.PATH_LENGTH = 100
+
+    def get_start_pose(self):
+        sim_scene: SimulationScene = SimulationScene(
+            env_name=self.version,
+            scene_name=f"sim_scene",
+            scene=self.scene,
+            dataset=self.dataset,
+            init_timestep=0,
+            freeze_agents=True,
+        )
+        obs = sim_scene.reset()
+        assert obs.agent_name[0] == 'ego', 'The first agent is not ego'
+        # We consider position of the first ego frame as origin
+        # This suppose is ok when the first frame front camera pose is set as origin
+        ego_start_pos = obs.curr_agent_state.position[0]
+        ego_start_yaw = obs.curr_agent_state.heading[0]
+        return ego_start_pos.numpy(), ego_start_yaw.item()
+
+    def xyzr_local2world(self, stat):
+        alpha = np.arctan(stat[0] / stat[1])
+        beta = self.ego_start_yaw - alpha
+        dist = np.linalg.norm(stat[:2])
+        delta_x = dist * np.cos(beta)
+        delta_y = dist * np.sin(beta)
+
+        world_stat = np.zeros(4)
+        world_stat[0] = delta_x + self.ego_start_pos[0]
+        world_stat[1] = delta_y + self.ego_start_pos[1]
+        world_stat[3] = stat[3] + self.ego_start_yaw
+
+        return world_stat
+
+    def batch_xyzr_world2local(self, stat):
+        beta = np.arctan((stat[:, 1] - self.ego_start_pos[1]) / (stat[:, 0] - self.ego_start_pos[0]))
+        alpha = self.ego_start_yaw - beta
+        dist = np.linalg.norm(stat[:, :2] - self.ego_start_pos, axis=1)
+        delta_x = dist * np.sin(alpha)
+        delta_y = dist * np.cos(alpha)
+
+        local_stat = np.zeros_like(stat)
+        local_stat[:, 0] = delta_x
+        local_stat[:, 1] = delta_y
+        local_stat[:, 3] = stat[:, 3] - self.ego_start_yaw
+
+        return local_stat
+
+    def get_route(self, stat):
+        # stat: a, b, height, yaw, v
+        curr_xyzr = self.xyzr_local2world(stat[:4].numpy())
+        
+        # lanes = self.vector_map.get_current_lane(curr_xyzr, max_dist=5, max_heading_error=np.pi/3)
+        lanes = self.vector_map.get_current_lane(curr_xyzr)
+
+        if len(lanes) > 0:
+            curr_lane = lanes[0]
+            path = self.batch_xyzr_world2local(curr_lane.center.xyzh)
+            total_path_length = np.linalg.norm(curr_lane.center.xy[1:] - curr_lane.center.xy[:-1], axis=1).sum()
+            # random select next lanes until reach PATH_LENGTH
+            while total_path_length < self.PATH_LENGTH:
+                next_lanes = list(curr_lane.next_lanes)
+                if len(next_lanes) == 0:
+                    break
+                next_lane = self.vector_map.get_road_lane(next_lanes[np.random.randint(len(next_lanes))])
+                path = np.vstack([path, self.batch_xyzr_world2local(next_lane.center.xyzh)])
+                total_path_length += np.linalg.norm(next_lane.center.xy[1:] - next_lane.center.xy[:-1], axis=1).sum()
+                curr_lane = next_lane
+        else:
+            path = None
+        return path

code/sim/utils/score_calculator.py

@@ -0,0 +1,562 @@
+import pickle
+from matplotlib import pyplot as plt
+from matplotlib.patches import Rectangle, Polygon
+from shapely.geometry import LineString, Point
+import numpy as np
+from shapely.geometry import Polygon as ShapelyPolygon
+from shapely.geometry import Point
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+import threading
+import argparse
+import os
+import open3d as o3d
+import torch
+from scipy.spatial.transform import Rotation as SCR
+
+ego_verts_canonic = np.array([[0.5, 0.5, 0], [0.5, -0.5, 0], [0.5, 0.5, 1.0], [0.5, -0.5, 1.0],
+                    [-0.5, -0.5, 0], [-0.5, 0.5, 0], [-0.5, -0.5, 1.0], [-0.5, 0.5, 1.0]])
+
+# Define boundaries
+boundaries = {
+    'max_abs_lat_accel': 4.89,  # [m/s^2]
+    'max_lon_accel': 2.40,  # [m/s^2]
+    'min_lon_accel': -4.05,  # [m/s^2]
+    'max_abs_yaw_accel': 1.93,  # [rad/s^2]
+    'max_abs_lon_jerk': 8.37,  # [m/s^3],
+    'max_abs_yaw_rate': 0.95,  # [rad/s]
+}
+
+score_weight = {
+    'ttc': 5,
+    'c': 2,
+    'ep': 5,
+}
+
+def create_rectangle(center_x, center_y, width, length, yaw):
+    """Create a rectangle polygon."""
+    cos_yaw = np.cos(yaw)
+    sin_yaw = np.sin(yaw)
+
+    x_offs = [length/2, length/2, -length/2, -length/2]
+    y_offs = [width/2, -width/2, -width/2, width/2]
+
+    x_pts = [center_x + x_off*cos_yaw - y_off *
+                sin_yaw for x_off, y_off in zip(x_offs, y_offs)]
+    y_pts = [center_y + x_off*sin_yaw + y_off *
+                cos_yaw for x_off, y_off in zip(x_offs, y_offs)]
+
+    return ShapelyPolygon(zip(x_pts, y_pts))
+
+def bg_collision_det(points, box):
+    O, A, B, C = box[0], box[1], box[2], box[5]
+    OA = A - O
+    OB = B - O
+    OC = C - O
+    POA, POB, POC = (points @ OA[..., None])[:, 0], (points @ OB[..., None])[:, 0], (points @ OC[..., None])[:, 0]
+    mask = (torch.dot(O, OA) < POA) & (POA < torch.dot(A, OA)) & \
+        (torch.dot(O, OB) < POB) & (POB < torch.dot(B, OB)) & \
+        (torch.dot(O, OC) < POC) & (POC < torch.dot(C, OC))
+    return True if torch.sum(mask) > 100 else False
+
+
+class ScoreCalculator:
+    def __init__(self, data):
+        self.data = data
+
+        self.pdms = 0.0
+        self.driving_score = None
+        pass
+
+    def transform_to_ego_frame(self, traj, ego_box):
+        """
+        Transform trajectory from global frame to ego-centric frame.
+
+        :param traj: List of tuples (x, y, yaw) in global frame
+        :param ego_box: Tuple (x, y, z, w, l, h, yaw) of ego vehicle in global frame
+        :return: Numpy array of transformed trajectory
+        """
+        ego_x, ego_y, _, _, _, _, ego_yaw = ego_box
+
+        # Create rotation matrix
+        c, s = np.cos(-ego_yaw), np.sin(-ego_yaw)
+        R = np.array([[c, -s], [s, c]])
+
+        # Transform each point
+        transformed_traj = []
+        for x, y, yaw in traj:
+            # Translate
+            x_translated, y_translated = x - ego_x, y - ego_y
+
+            # Rotate
+            x_rotated, y_rotated = R @ np.array([x_translated, y_translated])
+
+            # Adjust yaw
+            yaw_adjusted = yaw - ego_yaw
+
+            transformed_traj.append((x_rotated, y_rotated, yaw_adjusted))
+
+        return np.array(transformed_traj)
+
+    def get_vehicle_corners(self, x, y, yaw, length, width):
+        """
+        Calculate the corner points of the vehicle given its position, orientation, and dimensions.
+
+        :param x: x-coordinate of the vehicle's center
+        :param y: y-coordinate of the vehicle's center
+        :param yaw: orientation of the vehicle in radians
+        :param length: length of the vehicle
+        :param width: width of the vehicle
+        :return: numpy array of corner coordinates (4x2)
+        """
+        c, s = np.cos(yaw), np.sin(yaw)
+        front_left = np.array([x + c * length / 2 - s * width / 2,
+                               y + s * length / 2 + c * width / 2])
+        front_right = np.array([x + c * length / 2 + s * width / 2,
+                                y + s * length / 2 - c * width / 2])
+        rear_left = np.array([x - c * length / 2 - s * width / 2,
+                              y - s * length / 2 + c * width / 2])
+        rear_right = np.array([x - c * length / 2 + s * width / 2,
+                               y - s * length / 2 - c * width / 2])
+        return np.array([front_left, front_right, rear_right, rear_left])
+
+    def plot_trajectory_on_drivable_mask(self, drivable_mask, transformed_traj, vehicle_width, vehicle_length):
+        """
+        Plot the transformed trajectory and vehicle bounding boxes on the drivable mask.
+
+        :param drivable_mask: 2D numpy array representing the drivable area (200x200)
+        :param transformed_traj: Numpy array of transformed trajectory points
+        :param vehicle_width: Width of the vehicle in meters
+        :param vehicle_length: Length of the vehicle in meters
+        """
+
+        plt.figure(figsize=(10, 10))
+        plt.imshow(drivable_mask, cmap='gray', extent=[-50, 50, -50, 50])
+
+        # Scale factor (200 pixels represent 100 meters)
+        scale_factor = 200 / 100  # pixels per meter
+
+        # Plot trajectory
+        x_coords, y_coords, yaws = transformed_traj.T
+        plt.plot(x_coords, y_coords, 'r-', linewidth=2)
+
+        # Plot vehicle bounding boxes
+        for x, y, yaw in transformed_traj:
+            corners = self.get_vehicle_corners(
+                x, y, yaw, vehicle_length, vehicle_width)
+            plt.gca().add_patch(Polygon(corners, fill=False, edgecolor='blue'))
+
+        # Plot start and end points
+        plt.plot(x_coords[0], y_coords[0], 'go', markersize=10, label='Start')
+        plt.plot(x_coords[-1], y_coords[-1], 'bo', markersize=10, label='End')
+
+        plt.title('Trajectory and Vehicle Bounding Boxes on Drivable Mask')
+        plt.legend()
+        plt.xlabel('x (meters)')
+        plt.ylabel('y (meters)')
+        plt.grid(True)
+        plt.tight_layout()
+        plt.show()
+    
+    def _calculate_drivable_area_compliance(self, ground, traj, vehicle_width, vehicle_length):
+        m, n = 2, 2
+        dac = 1.0
+        for traj_i, (x, y, yaw) in enumerate(traj):
+            cnt = 0
+            c, s = np.cos(yaw), np.sin(yaw)
+            R = np.array([[c, -s], [s, c]])
+            ground_in_ego = (np.linalg.inv(R) @ (ground + np.array([-x, -y])).T).T
+            x_bins = np.linspace(-vehicle_length/2, vehicle_length/2, m+1)
+            y_bins = np.linspace(-vehicle_width/2, vehicle_width/2, n+1)
+            for xi in range(m):
+                for yi in range(n):
+                    min_x, max_x = x_bins[xi], x_bins[xi+1]
+                    min_y, max_y = y_bins[yi], y_bins[yi+1]
+                    ground_mask = (min_x < ground_in_ego[:, 0]) & (ground_in_ego[:, 0] < max_x) & \
+                                    (min_y < ground_in_ego[:, 1]) & (ground_in_ego[:, 1] < max_y)
+                    if ground_mask.sum() > 0:
+                        cnt += 1
+            drivable_ratio = cnt / (m*n)
+            if drivable_ratio < 0.3:
+                return 0
+            elif drivable_ratio < 0.5:
+                dac = 0.5
+        return dac
+
+    def _calculate_progress(self, planned_traj, ref_taj):
+        def calculate_curve_length(points):
+            """Calculate the total length of a curve given by a set of points."""
+            curve = LineString(points)
+            return curve.length
+
+        def project_curve_onto_curve(curve_a, curve_b):
+            """Project curve_b onto curve_a and calculate the projected length."""
+            projected_points = []
+            for point in curve_b.coords:
+                projected_point = curve_a.interpolate(
+                    curve_a.project(Point(point)))
+                projected_points.append(projected_point)
+            projected_curve = LineString(projected_points)
+            return projected_curve.length
+
+        # Create Shapely LineString objects
+        plan_curve = LineString([(x, y) for x, y, _ in planned_traj])
+        ref_curve = LineString([(x, y) for x, y, _ in ref_taj])
+
+        # Calculate lengths
+        plan_curve_length = calculate_curve_length(plan_curve)
+        ref_curve_length = calculate_curve_length(ref_curve)
+        projected_length = project_curve_onto_curve(ref_curve, plan_curve)
+        # print(f"plan_curve_length: {plan_curve_length}, ref_curve_length: {ref_curve_length}, project plan to ref_length: {projected_length}")
+
+        ep = 0.0
+        if max(plan_curve_length, ref_curve_length) < 5.0 or ref_curve_length < 1e-6:
+            ep = 1.0
+        else:
+            ep = projected_length / ref_curve_length
+        return ep
+
+    def _calculate_is_comfortable(self, traj, timestep):
+        """
+        Check if all kinematic parameters of a trajectory are within specified boundaries.
+
+        :param traj: List of tuples (x, y, yaw) representing the trajectory, in ego's local frame
+        :param timestep: Time interval between trajectory points in seconds
+        :return: 1.0 if all parameters are within boundaries, 0.0 otherwise
+        """
+
+        def calculate_trajectory_kinematics(traj, timestep):
+            """
+            Calculate kinematic parameters for a given trajectory.
+
+            :param traj: List of tuples (x, y, yaw) for each point in the trajectory
+            :param timestep: Time interval between each point in the trajectory
+            :return: Dictionary containing lists of calculated parameters
+            """
+            # Convert trajectory to numpy array for easier calculations
+            x, y, yaw = zip(*traj)
+            x, y, yaw = np.array(x), np.array(y), np.array(yaw)
+
+            # Calculate velocities
+            dx = np.diff(x) / timestep
+            dy = np.diff(y) / timestep
+
+            # Calculate yaw rate
+            dyaw = np.diff(yaw)
+            dyaw = np.where(dyaw > np.pi, dyaw - 2*np.pi, dyaw)
+            dyaw = np.where(dyaw < -np.pi, dyaw + 2*np.pi, dyaw)
+            dyaw = dyaw / timestep
+            ddyaw = np.diff(dyaw) / timestep
+
+            # Calculate speed
+            speed = np.sqrt(dx**2 + dy**2)
+
+            # Calculate accelerations
+            accel = np.diff(speed) / timestep
+            jerk = np.diff(accel) / timestep
+
+            # Calculate yaw rate (already calculated as dyaw)
+            yaw_rate = dyaw
+            # Calculate yaw acceleration
+            yaw_accel = ddyaw
+
+            lon_accel = accel
+            lat_accel = np.zeros_like(lon_accel)
+            lon_jerk = jerk
+
+            # Pad arrays to match the original trajectory length
+            yaw_rate = np.pad(yaw_rate, (0, 1), 'edge')
+            yaw_accel = np.pad(yaw_accel, (0, 2), 'edge')
+            lon_accel = np.pad(lon_accel, (0, 2), 'edge')
+            lat_accel = np.pad(lat_accel, (0, 2), 'edge')
+            lon_jerk = np.pad(lon_jerk, (0, 3), 'edge')
+
+            return {
+                'speed': speed,
+                'yaw_rate': yaw_rate,
+                'yaw_accel': yaw_accel,
+                'lon_accel': lon_accel,
+                'lat_accel': lat_accel,
+                'lon_jerk': lon_jerk,
+            }
+
+        # Calculate kinematic parameters
+        if len(traj) < 4:
+            return 1.0
+
+        kinematics = calculate_trajectory_kinematics(traj, timestep)
+
+        # Check each parameter against its boundary
+        checks = [
+            np.all(np.abs(kinematics['lat_accel']) <=
+                   boundaries['max_abs_lat_accel']),
+            np.all(kinematics['lon_accel'] <= boundaries['max_abs_lat_accel']),
+            np.all(kinematics['lon_accel'] >= boundaries['min_lon_accel']),
+            np.all(np.abs(kinematics['lon_jerk']) <=
+                   boundaries['max_abs_lon_jerk']),
+            np.all(np.abs(kinematics['yaw_accel']) <=
+                   boundaries['max_abs_yaw_accel']),
+            np.all(np.abs(kinematics['yaw_rate']) <=
+                   boundaries['max_abs_yaw_rate'])
+        ]
+
+        # if not all(checks):
+        #     print(traj)
+        #     print(kinematics)
+        print(f"comfortable: {checks}")
+
+        # Return 1.0 if all checks pass, 0.0 otherwise
+        return 1.0 if all(checks) else 0.0
+
+    def _calculate_no_collision(self, ego_box, planned_traj, obs_lists, scene_xyz):
+        ego_x, ego_y, z, ego_w, ego_l, ego_h, ego_yaw = ego_box
+        ego_verts_local = ego_verts_canonic * np.array([ego_l, ego_w, ego_h])
+        for idx in range(planned_traj.shape[0]):
+            ego_x, ego_y, ego_yaw = planned_traj[idx]  # ego_state= (x,y,yaw)
+            ego_trans_mat = np.eye(4)
+            ego_trans_mat[:3, :3] = SCR.from_euler('z', ego_yaw).as_matrix()
+            ego_trans_mat[:3, 3] = np.array([ego_x, ego_y, z])
+            ego_verts_global = (ego_trans_mat[:3, :3] @ ego_verts_local.T).T + ego_trans_mat[:3, 3]
+            ego_verts_global = torch.from_numpy(ego_verts_global).float().cuda()
+            bk_collision = bg_collision_det(scene_xyz, ego_verts_global)
+            # scene_local = scene_xyz - np.array([ego_x, ego_y, z])
+            # bk_collision = np.sum(
+            #     (-ego_l/2 < scene_local[:, 0]) & (scene_local[:, 0] < ego_l/2) & \
+            #     (-ego_w/2 < scene_local[:, 1]) & (scene_local[:, 1] < ego_w/2) & \
+            #     (-ego_h/2 < scene_local[:, 2]) & (scene_local[:, 2] < ego_h/2)
+            # ) > 100
+            if bk_collision:
+                print(f"collision with background detected! @ timestep{idx}")
+                return 0.0
+            ego_poly = create_rectangle(ego_x, ego_y, ego_w, ego_l, ego_yaw)
+            obs_list = obs_lists[idx if idx < len(obs_lists) else -1]
+            for obs in obs_list:
+                # obs = (x,y,z,w,l,h,yaw)
+                obs_x, obs_y, _, obs_w, obs_l, _, obs_yaw = obs
+                obs_poly = create_rectangle(
+                    obs_x, obs_y, obs_w, obs_l, obs_yaw)
+                if ego_poly.intersects(obs_poly):
+                    print(f"collision with obstacle detected! @ timestep{idx}")
+                    print(
+                        f"ego_poly: {(ego_x, ego_y, ego_yaw,obs_w, obs_l)}, obs_poly: {(obs_x, obs_y, obs_yaw,obs_w, obs_l )}")
+                    return 0.0  # Collision detected
+        return 1.0
+
+    def _calculate_time_to_collision(self, ego_box, planned_traj, obs_lists, scene_xyz, timestep):
+        # breakpoint()
+        t_list = [0.5, 1]  # ttc time
+
+        for t in t_list:
+            # Calculate velocities
+            velocities = np.diff(planned_traj[:, :2], axis=0) / timestep
+
+            # Use the velocity of the second point for the first point
+            velocities = np.vstack([velocities[0], velocities])
+
+            # Calculate the displacement
+            displacement = velocities * t
+
+            # Create the new trajectory
+            new_traj = planned_traj.copy()
+            new_traj[:, :2] += displacement
+
+            is_collide_score = self._calculate_no_collision(
+                ego_box, new_traj, obs_lists, scene_xyz)
+            if is_collide_score == 0.0:
+                print(f" failed to pass ttc collision check, t={t}")
+                # breakpoint()
+                return 0.0
+
+        return 1.0
+
+    def calculate(self, ):
+
+        print(f"current exp has {len(self.data['frames'])} frames")
+        if len(self.data['frames']) == 0:
+            return None
+        # todo: time_step need modify
+        score_list = {}
+        for i in range(0, len(self.data['frames']), 1):
+            frame = self.data['frames'][i]
+            if frame['is_key_frame'] == False:
+                continue
+
+            print(f"frame {i} / {len(self.data['frames'])}")
+            timestamp = frame['time_stamp']
+            planned_last_timestamp = timestamp + \
+                len(frame['planned_traj']['traj']) * \
+                frame['planned_traj']['timestep']
+            ego_x, ego_y, _, ego_w, ego_l, _, ego_yaw = frame['ego_box']
+            # frame['planned_traj']['traj']
+            if len(frame['planned_traj']['traj'])<2:
+                continue
+            traj = frame['planned_traj']['traj']
+
+            planned_traj = np.concatenate(([np.array([ego_x, ego_y, ego_yaw])], traj), axis=0)
+            # print(planned_traj)
+
+            # if the car is stopped, there may be error in the yaw of the planned trajectory
+            traj_distance =  np.linalg.norm(planned_traj[-1, :2]  - planned_traj[0, :2] )
+            if traj_distance<1:
+                 planned_traj[:, 2] = planned_traj[0, 2]  # set all yaw to the first yaw
+
+            current_timestamp = timestamp
+            current_frame_idx = i
+            obs_lists = []
+            while current_timestamp <= planned_last_timestamp+1e-5:
+                if abs(current_timestamp - self.data['frames'][current_frame_idx]['time_stamp']) < 1e-5:
+                    obs_list = []
+                    for idx, obj in enumerate(self.data['frames'][current_frame_idx]['obj_boxes']):
+                        # obs_list.append(obj)
+                        if self.data['frames'][current_frame_idx]['obj_names'][idx] == 'car':
+                            obs_list.append(obj)
+                    obs_lists.append(obs_list)
+                    current_timestamp += frame['planned_traj']['timestep']
+
+                current_frame_idx += 1
+                if current_frame_idx >= len(self.data['frames']):
+                    break
+
+            # breakpoint()
+            # plt.imshow(frame['drivable_mask'].astype(np.uint8))
+            # plt.show()
+            # transformed_traj = self.transform_to_ego_frame(frame['planned_traj']['traj'], frame['ego_box'])
+            transformed_traj = self.transform_to_ego_frame(
+                planned_traj, frame['ego_box'])
+            # breakpoint()
+
+            score_nc = self._calculate_no_collision(
+                frame['ego_box'], planned_traj, obs_lists, self.data['scene_xyz'])
+            # score_nc = 0.0 if frame['collision'] else 1.0
+            score_dac = self._calculate_drivable_area_compliance(
+                self.data['ground_xy'], planned_traj, ego_w, ego_l)
+            score_ttc = self._calculate_time_to_collision(
+                frame['ego_box'], planned_traj, obs_lists, self.data['scene_xyz'], frame['planned_traj']['timestep'])
+            score_c = self._calculate_is_comfortable(
+                transformed_traj, frame['planned_traj']['timestep'])
+            # score_ep = self._calculate_progress(
+            #     planned_traj, ref_traj)
+
+            score_pdms = score_nc*score_dac*(score_weight['ttc']*score_ttc+score_weight['c']*score_c)/(
+                score_weight['ttc']+score_weight['c'])
+            print('nc, dac, ttc, com, pdms', [score_nc, score_dac, score_ttc, score_c, score_pdms])
+            score_list[timestamp] = {'nc': score_nc, 'dac': score_dac,
+                                     'ttc': score_ttc, 'c': score_c, 'pdms': score_pdms}
+        
+        totals = {metric: 0 for metric in next(iter(score_list.values()))}
+        for scores in score_list.values():
+            for metric, value in scores.items():
+                totals[metric] += value
+
+        # avg scores
+        num_entries = len(score_list)
+        averages = {metric: total / num_entries for metric,
+                    total in totals.items()}
+
+        #     writer.writerow(averages.values())
+
+        mean_score = averages['pdms']
+        route_completion = max([f['rc'] for f in self.data['frames']])
+        route_completion = route_completion if route_completion < 1 else 1.0
+        driving_score = mean_score*route_completion
+        return mean_score, route_completion, driving_score, averages
+
+
+def calculate(data):
+    print(f"this pkl file contains {len(data)} experiment records.")
+    # print(f"the first item metadata is {data[0]['metas']}.")
+    # breakpoint()
+
+    def process_exp_data(exp_data):
+        score_calc = ScoreCalculator(exp_data)
+        score = score_calc.calculate()
+        print(f"The score of experiment is {score}.")
+        final_score_dict = score[3]
+        final_score_dict['rc'] = score[1]
+        final_score_dict['hdscore'] = score[2]
+        return final_score_dict
+
+    def multi_threaded_process(data, max_workers=None):
+        all_averages = []
+
+        # Using thread locks for thread-safe append operations
+        lock = threading.Lock()
+
+        def append_result(future):
+            result = future.result()
+            with lock:
+                all_averages.append(result)
+
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            futures = [executor.submit(process_exp_data, exp_data)
+                       for exp_data in data]
+            for future in futures:
+                future.add_done_callback(append_result)
+
+        return all_averages
+
+    all_averages = multi_threaded_process(data)
+
+    collected_values = defaultdict(list)
+    for averages in all_averages:
+        for key, value in averages.items():
+            collected_values[key].append(value)
+
+    # Calculation of mean and standard deviation for each indicator
+    results = {}
+    for key, values in collected_values.items():
+        avg = np.mean(values)
+        # std = np.std(values)
+        results[key] = f"{avg:.4f}"
+
+    # Output Results
+    print("=============================Results=============================")
+    for key, value in results.items():
+        print(f"'{key}': {value}")
+    return results
+
+def parse_data(test_path):
+    data_file_name = os.path.join(test_path, "data.pkl")
+    ground_pcd_file_name = os.path.join(test_path, "ground.ply")
+    scene_pcd_file_name = os.path.join(test_path, "scene.ply")
+
+    # Open the file and load the data
+    with open(data_file_name, 'rb') as f:
+        data = pickle.load(f)
+    
+    ground_pcd = o3d.io.read_point_cloud(ground_pcd_file_name)
+    ground_xyz = np.asarray(ground_pcd.points) # in camera coordinates
+    ground_xy = np.stack([ground_xyz[:, 2], -ground_xyz[:, 0]], axis=1) # in imu coordinates
+    
+    scene_pcd = o3d.io.read_point_cloud(scene_pcd_file_name)
+    scene_xyz = np.asarray(scene_pcd.points) # in camera coordinates
+    # in imu coordinates
+    scene_xyz = np.stack([scene_xyz[:, 2], -scene_xyz[:, 0], -scene_xyz[:, 1]], axis=1) 
+    
+    data[0]['ground_xy'] = ground_xy
+    data[0]['scene_xyz'] = torch.from_numpy(scene_xyz).cuda()
+    # data[0]['scene_xyz'] = scene_xyz
+    return data
+
+
+def hugsim_evaluate(test_data, ground_xyz, scene_xyz):
+    ground_xy = np.stack([ground_xyz[:, 2], -ground_xyz[:, 0]], axis=1) # in imu coordinates
+    scene_xyz = np.stack([scene_xyz[:, 2], -scene_xyz[:, 0], -scene_xyz[:, 1]], axis=1) 
+    test_data[0]['ground_xy'] = ground_xy
+    test_data[0]['scene_xyz'] = torch.from_numpy(scene_xyz).float().cuda()
+    results = calculate(test_data)
+    return results
+
+
+def get_opts():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--test_path', type=str, required=True)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = get_opts()
+    data = parse_data(args.test_path)
+    
+    # Call the main function with the loaded data
+    calculate(data)

code/sim/utils/sim_utils.py

@@ -0,0 +1,122 @@
+import numpy as np
+from scipy.spatial.transform import Rotation as SCR
+import math
+from scene.cameras import Camera
+from sim.ilqr.lqr import plan2control
+from omegaconf import OmegaConf
+
+def rt2pose(r, t, degrees=False):
+    pose = np.eye(4)
+    pose[:3, :3] = SCR.from_euler('XYZ', r, degrees=degrees).as_matrix()
+    pose[:3, 3] = t
+    return pose
+
+def pose2rt(pose, degrees=False):
+    r = SCR.from_matrix(pose[:3, :3]).as_euler('XYZ', degrees=degrees)
+    t = pose[:3, 3]
+    return r, t
+    
+def load_camera_cfg(cfg):
+    cam_params = {}
+    cams = OmegaConf.to_container(cfg.cams, resolve=True)
+    for cam_name, cam in cams.items():
+        v2c = rt2pose(cam['extrinsics']['v2c_rot'], cam['extrinsics']['v2c_trans'], degrees=True)
+        l2c = rt2pose(cam['extrinsics']['l2c_rot'], cam['extrinsics']['l2c_trans'], degrees=True)
+        cam_intrin = cam['intrinsics']
+        cam_intrin['fovx'] = cam_intrin['fovx'] / 180.0 * np.pi
+        cam_intrin['fovy'] = cam_intrin['fovy'] / 180.0 * np.pi
+        cam_params[cam_name] = {'intrinsic': cam_intrin, 'v2c': v2c, 'l2c': l2c}
+        
+    rect_mat = np.eye(4)
+    if 'cam_rect' in cfg:
+        rect_mat[:3, :3] = SCR.from_euler('XYZ', cfg.cam_rect.rot, degrees=True).as_matrix()
+        rect_mat[:3, 3] = np.array(cfg.cam_rect.trans)
+        
+    return cam_params, OmegaConf.to_container(cfg.cam_align, resolve=True), rect_mat
+
+def fov2focal(fov, pixels):
+    return pixels / (2 * math.tan(fov / 2))
+
+def focal2fov(focal, pixels):
+    return 2*math.atan(pixels/(2*focal))
+
+def create_cam(intrinsic, c2w):
+    fovx, fovy = intrinsic['fovx'], intrinsic['fovy']
+    h, w = intrinsic['H'], intrinsic['W']
+    K = np.eye(4)
+    K[0, 0], K[1, 1] = fov2focal(fovx, w), fov2focal(fovy, h)
+    K[0, 2], K[1, 2] = intrinsic['cx'], intrinsic['cy']
+    cam = Camera(K=K, c2w=c2w, width=w, height=h,
+                image=np.zeros((h, w, 3)), image_name='', dynamics={})
+    return cam
+
+def traj2control(plan_traj, info):
+    """
+        The input plan trajectory is under lidar coordinates
+        x to right, y to forward and z to upward
+    """
+    plan_traj_stats = np.zeros((plan_traj.shape[0]+1, 5))
+    plan_traj_stats[1:, :2] = plan_traj[:, [1,0]]
+    prev_a, prev_b = 0, 0
+    for i, (a, b) in enumerate(plan_traj):
+        rot = np.arctan((b - prev_b)/(a - prev_a))
+        plan_traj_stats[i+1, 2] = rot
+    curr_stat = np.array(
+        [0, 0, 0, info['ego_velo'], info['ego_steer']]
+    )
+    acc, steer_rate = plan2control(plan_traj_stats, curr_stat)
+    return acc, steer_rate
+
+def dense_cam_poses(cam_poses, cmds):
+    
+    for i in range(5):
+        dense_poses = []
+        dense_cmds = []
+        for i in range(cam_poses.shape[0]-1):
+            cam1 = cam_poses[i]
+            cam2 = cam_poses[i+1]
+            dense_poses.append(cam1)
+            dense_cmds.append(cmds[i])
+            if np.linalg.norm(cam1[:3, 3]-cam2[:3, 3]) > 0.1:
+                euler1 = SCR.from_matrix(cam1[:3, :3]).as_euler("XYZ")
+                euler2 = SCR.from_matrix(cam2[:3, :3]).as_euler("XYZ")
+                interp_euler = (euler1 + euler2) / 2
+                interp_trans = (cam1[:3, 3] + cam2[:3, 3]) / 2
+                interp_pose = np.eye(4)
+                interp_pose[:3, :3] = SCR.from_euler("XYZ", interp_euler).as_matrix()
+                interp_pose[:3, 3] = interp_trans
+                dense_poses.append(interp_pose)
+                dense_cmds.append(cmds[i])
+        dense_poses.append(cam_poses[-1])
+        dense_poses = np.stack(dense_poses)
+        cam_poses = dense_poses
+        cmds = dense_cmds
+        
+    return cam_poses, cmds
+
+def traj_transform_to_global(traj, ego_box):
+        """
+        Transform trajectory from ego-centeric frame to global frame
+        """
+        ego_x, ego_y, _, _, _, _, ego_yaw = ego_box
+        global_points = [
+            (
+                ego_x
+                + px * math.cos(ego_yaw)
+                - py * math.sin(ego_yaw),
+                ego_y
+                + px * math.sin(ego_yaw)
+                + py * math.cos(ego_yaw),
+            )
+            for px, py in traj
+        ]
+        global_trajs = []
+        for i in range(1, len(global_points)):
+            x1, y1 = global_points[i - 1]
+            x2, y2 = global_points[i]
+            dx, dy = x2 - x1, y2 - y1
+            # distance = math.sqrt(dx**2 + dy**2)
+            yaw = math.atan2(dy, dx)
+            global_trajs.append((x1, y1, yaw))
+        return global_trajs
+        

code/submodules/Pplan/Policy/base.py

@@ -0,0 +1,16 @@
+import abc
+
+
+class Policy(abc.ABC):
+    def __init__(self, device, *args, **kwargs):
+        self.device = device
+
+    @abc.abstractmethod
+    def get_action(self, obs_dict, **kwargs):
+        """Predict an action based on the input observation """
+        pass
+
+    @abc.abstractmethod
+    def eval(self):
+        """Set the policy to evaluation mode"""
+        pass

code/submodules/Pplan/Sampling/forward_sampler.py

@@ -0,0 +1,141 @@
+from logging import raiseExceptions
+import numpy as np
+import torch
+import pdb
+from ..utils import geometry_utils as GeoUtils
+import matplotlib.pyplot as plt
+from scipy.interpolate import interp1d
+import random
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Final,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
+class ForwardSampler(object):
+    def __init__(self,dt:float, acce_grid:list,dhm_grid:list, dhf_grid:list,  max_rvel=8,max_steer=0.5, vbound=[-5.0, 30],device="cuda" if torch.cuda.is_available() else "cpu"):
+        self.device = device
+        self.accels = torch.tensor(acce_grid,device=self.device)
+        self.dhf_grid = torch.tensor(dhf_grid,device=self.device)
+        self.dhm_grid = torch.tensor(dhm_grid,device=self.device)
+        self.max_rvel = max_rvel
+        self.vbound = vbound
+        self.max_steer = max_steer
+        self.dt = dt
+
+        
+    def velocity_plan(self,x0:torch.Tensor,T:int,acce:Optional[torch.Tensor]=None):
+        """plan velocity profile
+
+        Args:
+            x0 (torch.Tensor): [B, 4], X,Y,v,heading
+            T (int): time horizon
+            acce (torch.Tensor): [B, N]
+        """
+        bs = x0.shape[0]
+        if acce is None:
+            acce = self.accels[None,:].repeat_interleave(bs,0)
+        v0 = x0[...,2] # [B]
+        vdes = v0[:,None,None]+torch.arange(T,device=self.device)[None,None]*acce[:,:,None]*self.dt
+        vplan = torch.clip(vdes,min=self.vbound[0],max=self.vbound[1])
+        return vplan # [B, N, T]
+
+    def lateral_plan(self,x0:torch.Tensor,vplan:torch.Tensor,dhf:torch.Tensor,dhm:torch.Tensor,T:int,bangbang=True):
+        """plan lateral profile, 
+            steering plan that ends with the desired heading change with mean heading change equal to dhm, if feasible
+        Args:
+            x0 (torch.Tensor): [B, 4], X,Y,v,heading
+            vplan (torch.Tensor): [B, N, T] velocity profile 
+            dhf (torch.Tensor): [B, M] desired heading change at the end of the horizon
+            dhm (torch.Tensor): [B, M] mean heading change at the end of the horizon
+            T (int): horizon
+        """
+        # using a linear steering profile
+        bs,M = dhf.shape
+        N = vplan.shape[1]
+        vplan = vplan[:,:,None] # [B, N, 1, T]
+        vl = torch.cat([x0[:,2].reshape(-1,1,1,1).repeat_interleave(N,1),vplan[...,:-1]],-1)
+        acce = vplan-vl
+        
+        c0 = torch.abs(vl)
+        c1 = torch.cumsum(c0*self.dt,-1)
+        c2 = torch.cumsum(c1*self.dt,-1)
+        c3 = torch.cumsum(c2*self.dt,-1)
+                
+        
+        # algebraic equation: c1[T]*a0+c2[T]*a1 = dhf, c2[T]*a0+c3[T]*a1 = dhm
+        
+        a0 = (c3[...,-1]*dhf.unsqueeze(1)-c2[...,-1]*dhm.unsqueeze(1))/(c1[...,-1]*c3[...,-1]-c2[...,-1]**2) # [B, N, M]
+        a1 = (dhf.unsqueeze(1)-c1[...,-1]*a0)/c2[...,-1]
+        
+        yawrate = a0[...,None]*c0+a1[...,None]*c1
+
+        if bangbang:
+        # turn into bang-bang control to reduce the peak steering value, but the mean heading value is not retained
+            pos_flag = (yawrate>0)
+            neg_flag = ~pos_flag
+            mean_pos_steering = (yawrate*pos_flag).sum(-1)/((c0*pos_flag).sum(-1)+1e-6)
+            mean_neg_steering = (yawrate*neg_flag).sum(-1)/((c0*neg_flag).sum(-1)+1e-6)
+            mean_pos_steering = torch.clip(mean_pos_steering,min=-self.max_steer,max=self.max_steer)
+            mean_neg_steering = torch.clip(mean_neg_steering,min=-self.max_steer,max=self.max_steer)
+            bb_yawrate = (mean_pos_steering[...,None]*pos_flag+mean_neg_steering[...,None]*neg_flag)*c0
+            bb_yawrate = torch.clip(bb_yawrate,min=-self.max_rvel/c0,max=self.max_rvel/c0)
+            dh = torch.cumsum(bb_yawrate*self.dt,-1)
+        else:
+            yawrate = torch.clip(yawrate,min=-self.max_rvel/c0,max=self.max_rvel/c0)
+            yawrate = torch.clip(yawrate,min=-self.max_steer*c0,max=self.max_steer*c0)
+            dh = torch.cumsum(yawrate*self.dt,-1)
+        heading = x0[...,3,None,None,None]+dh
+        
+        vx = vplan*torch.cos(heading)
+        vy = vplan*torch.sin(heading)
+        traj = torch.stack([x0[:,None,None,None,0]+vx.cumsum(-1)*self.dt,
+                            x0[:,None,None,None,1]+vy.cumsum(-1)*self.dt,
+                            vplan.repeat_interleave(M,2),
+                            heading],-1)
+        t = torch.arange(1,T+1,device=self.device)[None,None,None,:,None].repeat(bs,N,M,1,1)*self.dt
+        xyvaqrt = torch.cat([traj[...,:3],acce[...,None].repeat_interleave(M,2),traj[...,3:],yawrate[...,None],t],-1)
+        return xyvaqrt.reshape(bs,N*M,T,-1) # [B, N*M, T, 7]
+        
+    def sample_trajs(self,x0,T,bangbang=True):
+        # velocity sample
+        vplan = self.velocity_plan(x0,T)
+        bs = x0.shape[0]
+        dhf = self.dhf_grid
+        dhm = self.dhm_grid
+        Mf = dhf.shape[0]
+        Mm = dhm.shape[0]
+        dhm = dhm.repeat(Mf).unsqueeze(0).repeat_interleave(bs,0)
+        dhf = dhf.repeat_interleave(Mm,0).unsqueeze(0).repeat_interleave(bs,0)+dhm
+        return self.lateral_plan(x0,vplan,dhf,dhm,T,bangbang)
+        
+        
+        
+def test():
+    sampler = ForwardSampler(acce_grid=[-4,-2,0,2,4],dhm_grid=torch.linspace(-0.7,0.7,9),dhf_grid=[-0.4,0,0.4],dt=0.1)
+    x0 = torch.tensor([0,0,1.,0.],device="cuda").unsqueeze(0).repeat_interleave(3,0)
+    T = 10
+    # vel_grid = sampler.velocity_plan(x0,T)
+    # dhf = torch.tensor([0.5,0,-0.5]).repeat(3).unsqueeze(0)
+    # dhm = torch.tensor([0.2,0,-0.2]).repeat_interleave(3,0).unsqueeze(0)
+    traj = sampler.sample_trajs(x0,T,bangbang=False)
+    traj = traj[0].reshape(-1,T,7).cpu().numpy()
+    import matplotlib.pyplot as plt
+    fig,ax = plt.subplots()
+    for i in range(traj.shape[0]):
+        ax.plot(traj[i,:,0],traj[i,:,1])
+    plt.show()
+    
+    
+        
+if __name__ == "__main__":
+    test()
+        
+        
+