PIN-SLAM 코드 읽기 - 1. Contribution과 전체 Pipeline

들어가며

요즘 PIN-SLAM 코드를 읽고 있다. 이 글은 남에게 완성된 tutorial을 제공한다기보다는, 내가 코드를 읽고 이해하기 위해 남기는 공부 노트에 가깝다. 다만 나중에 처음 보는 사람도 따라올 수 있도록, 최대한 block-wise로 설명해보려고 한다.

PIN-SLAM의 논문 제목은 다음과 같다.

LiDAR SLAM Using a Point-Based Implicit Neural Representation for Achieving Global Map Consistency

제목이 거의 모든 것을 말해준다. 핵심 단어는 세 개이다.

LiDAR SLAM: 입력은 LiDAR point cloud이고, trajectory와 map을 동시에 추정한다.
Point-Based Implicit Neural Representation: map을 단순 point cloud나 voxel grid가 아니라, neural point 기반의 implicit SDF field로 표현한다.
Global Map Consistency: loop closure와 pose graph optimization 이후, map 자체도 같이 변형해서 global consistency를 맞춘다.

이 시리즈는 PIN-SLAM을 논문 수식 위주로 보는 글이 아니다. 코드를 읽는 순서대로, “이 block은 왜 있고, 어떤 데이터를 받아 어떤 데이터를 내보내는가”를 이해하는 것이 목표이다.

PIN-SLAM의 contribution을 한 문장으로

내가 이해한 PIN-SLAM의 contribution은 다음 한 문장으로 요약할 수 있다.

LiDAR scan을 online으로 처리하면서, sparse neural point들이 만드는 implicit SDF map을 학습하고, 그 map에 scan을 직접 registration하며, loop closure가 생기면 pose graph와 neural point map을 함께 고쳐 globally consistent한 mesh까지 뽑는 SLAM system.

조금 풀어 쓰면 다음 네 가지가 중요하다.

1. Map이 point cloud가 아니라 implicit SDF field이다

일반적인 LiDAR SLAM에서는 map이 보통 point cloud, voxel, surfel, feature map 같은 형태이다. PIN-SLAM에서는 map이 SDF를 query할 수 있는 함수에 가깝다.

즉 어떤 3D point \(p\)를 넣으면,

\[S(p) \approx \text{surface까지의 signed distance}\]

를 뱉는 map을 만든다. 이 SDF를 이용하면 두 가지가 가능해진다.

scan point가 현재 map surface 위에 잘 올라가는지 residual로 볼 수 있다.
공간 전체를 query해서 marching cubes로 mesh를 만들 수 있다.

2. SDF field를 dense voxel grid가 아니라 sparse neural point로 표현한다

PIN-SLAM은 모든 공간을 dense하게 저장하지 않는다. 대신 관측된 surface 근처에 neural point를 둔다. 각 neural point는 대략 이런 정보를 가진다.

위치 \(x_i\)
orientation \(q_i\)
geometric latent feature \(f_i\)
생성/갱신 timestamp
certainty

query point \(p\)가 들어오면 주변 neural point들을 찾고, 각 neural point의 feature와 local coordinate를 MLP decoder에 넣어 SDF를 예측한다. 이 구조 덕분에 map은 compact하고, 관측된 영역 위주로만 자란다.

3. Odometry는 point-to-implicit registration이다

PIN-SLAM의 tracking은 scan-to-scan이나 scan-to-point-map ICP가 아니다. 현재 scan point들을 현재 추정 pose로 world에 올린 뒤, 그 point들의 SDF residual을 줄이는 pose update를 구한다.

즉 대략 이런 문제를 푼다.

\[\min_T \sum_i S(Tp_i)^2\]

여기서 \(p_i\)는 현재 scan의 sensor frame point이고, \(T \in SE(3)\)는 이 point를 map/world coordinate로 올리는 현재 pose이다. 즉 \(Tp_i\)는 homogeneous coordinate transform을 줄여 쓴 표기이고, \(S(\cdot)\)는 현재 neural point map이 표현하는 implicit SDF이다. 코드에서는 utils/tracker.py의 Tracker.tracking()과 implicit_reg()가 이 역할을 한다.

4. Loop closure 이후 map도 같이 변형한다

보통 pose graph optimization을 하면 pose들이 update된다. 그로 인해 이미 만들어둔 map을 어떻게 고칠지가 문제가 된다. PIN-SLAM은 neural point들이 timestamp를 가지고 있기 때문에, PGO 전후 pose 차이를 각 neural point에 적용해서 map을 같이 변형한다.

이게 제목의 global map consistency와 직접 연결된다. loop closure 이후 trajectory만 맞는 것이 아니라, 중복되거나 어긋난 implicit map도 같이 정렬되는 것이다.

코드 전체 구조

로컬 기준으로 내가 읽은 코드는 /Users/fudxo/pearl-slam/PIN_SLAM에 있다. 핵심 파일은 다음 정도이다.

파일	역할
`pin_slam.py`	전체 SLAM main loop
`dataset/slam_dataset.py`	frame loading, preprocessing, pose bookkeeping
`model/neural_points.py`	sparse neural point map, voxel hash, local map, feature query
`model/decoder.py`	latent feature + local coordinate를 SDF/semantic/color로 decode하는 MLP
`utils/data_sampler.py`	LiDAR ray에서 SDF training sample 생성
`utils/mapper.py`	neural map update, online SDF training, optional BA
`utils/tracker.py`	point-to-implicit odometry / loop registration
`utils/loop_detector.py`	neural point map context 또는 scan context 기반 loop candidate 탐지
`utils/pgo.py`	GTSAM 기반 pose graph optimization
`utils/mesher.py`	SDF grid query + marching cubes mesh reconstruction

처음 읽을 때는 pin_slam.py만 잡고 가면 된다. 나머지 파일은 pin_slam.py가 부르는 부품이다.

main loop 한눈에 보기

pin_slam.py의 중심은 아래 loop이다.

for frame_id in tqdm(range(dataset.total_pc_count)):
    # I. Load data and preprocessing
    # II. Odometry
    # III. Loop detection and PGO
    # IV. Mapping and bundle adjustment
    # V. Mesh reconstruction and visualization

# VI. Save results

이 구조만 머리에 넣어두면 PIN-SLAM이 갑자기 덜 복잡해진다. 결국 매 frame마다 하는 일은 다음과 같다.

scan을 읽고 전처리한다.
현재 map에 scan을 registration해서 현재 pose를 추정한다.
loop closure 후보를 찾고, 필요하면 pose graph를 최적화한다.
현재 scan으로 neural point map을 업데이트하고 SDF를 학습한다.
필요하면 mesh/SDF slice/trajectory를 visualization한다.
마지막에 trajectory, neural point map, mesh를 저장한다.

이제 block별로 조금 더 자세히 보자.

초기화 block

run_pin_slam()이 시작되면 먼저 config를 읽는다.

config = Config()
config.load(config_path)

그 뒤 핵심 object들을 만든다.

geo_mlp = Decoder(...)
sem_mlp = Decoder(...) if config.semantic_on else None
color_mlp = Decoder(...) if config.color_on else None

neural_points = NeuralPoints(config)
dataset = SLAMDataset(config)
tracker = Tracker(config, neural_points, mlp_dict)
mapper = Mapper(config, dataset, neural_points, mlp_dict)
mesher = Mesher(config, neural_points, mlp_dict)
pgm = PoseGraphManager(config)

여기서 중요한 점은 neural_points와 mlp_dict가 여러 module에 공유된다는 것이다.

Tracker는 현재 neural map을 query해서 pose를 추정한다.
Mapper는 같은 neural map과 decoder를 학습시킨다.
Mesher는 같은 neural map을 query해서 mesh를 만든다.

즉 PIN-SLAM에서 map은 NeuralPoints + Decoder 조합이다. 둘 중 하나만 있어서는 SDF를 만들 수 없다.

Block I. Load data and preprocessing

main loop의 첫 block은 data loading이다.

if config.use_dataloader:
    dataset.read_frame_with_loader(frame_id)
else:
    dataset.read_frame(frame_id)

valid_frame = dataset.preprocess_frame()

여기서 point cloud를 읽고, range filtering, downsampling, pose guess, color/semantic label 준비 같은 일을 한다. 이 글에서는 dataset 내부까지 깊게 들어가지는 않는다. PIN-SLAM의 핵심 contribution은 dataset loader가 아니라 neural implicit map과 SLAM loop에 있기 때문이다.

초보자 입장에서는 이렇게 생각하면 된다.

이 block의 주요 출력은 현재 frame 데이터이다.

dataset.cur_source_points: tracking에 사용할 source point들.
dataset.cur_point_cloud_torch: mapping에 사용할 현재 point cloud.
dataset.cur_pose_guess_torch: tracker가 시작할 초기 pose guess.

Block II. Odometry

두 번째 block은 현재 frame pose를 추정한다.

tracking_result = tracker.tracking(
    dataset.cur_source_points,
    dataset.cur_pose_guess_torch,
    dataset.cur_source_colors,
    dataset.cur_source_normals,
    vis_result=config.o3d_vis_on,
)
cur_pose_torch, cur_odom_cov, weight_pc_o3d, valid_flag = tracking_result
dataset.update_odom_pose(cur_pose_torch)

dataset.cur_pose_guess_torch는 constant velocity 같은 motion model로 만든 초기값이다. Tracker.tracking()은 이 초기 pose를 시작점으로 해서, scan point들이 implicit SDF surface에 잘 맞도록 pose를 반복적으로 업데이트한다.

여기서 중요한 mental model은 이것이다.

현재 scan을 map에 붙이는 기준은 nearest neighbor point-to-point distance가 아니라, neural map이 예측하는 SDF value이다.

이 부분은 3편에서 자세히 본다.

Block III. Loop detection and PGO

세 번째 block은 SLAM backend이다. config.pgo_on이 켜져 있으면 pose graph에 frame node와 odometry factor를 추가한다.

pgm.add_frame_node(frame_id, dataset.pgo_poses[frame_id])
pgm.add_odometry_factor(frame_id, frame_id-1, dataset.last_odom_tran)

그 다음 loop candidate를 찾는다. PIN-SLAM에는 크게 두 종류가 있다.

local loop: 현재 pose가 과거 pose 근처에 다시 왔는지 거리로 확인한다.
global loop: scan context 또는 neural point map context descriptor로 과거 장소를 찾는다.

후보가 나오면 바로 pose graph에 넣지 않는다. 먼저 Tracker.tracking(..., loop_reg=True)로 loop transform을 refine한다. 즉 loop closure도 결국 현재 scan을 과거 local map에 point-to-implicit registration하는 식으로 검증한다.

성공하면 pose graph를 optimize한다.

pgm.optimize_pose_graph()
pose_diff_torch = torch.tensor(pgm.get_pose_diff(), ...)
neural_points.adjust_map(pose_diff_torch)
neural_points.recreate_hash(...)
mapper.transform_data_pool(pose_diff_torch)
dataset.update_poses_after_pgo(pgm.pgo_poses)

여기서 PIN-SLAM의 중요한 부분이 나온다. PGO 이후 pose만 바꾸는 것이 아니라, neural point map과 mapper의 data pool도 같이 변형한다. 그래야 이후 tracking/mapping이 새 pose graph 기준으로 이어진다.

Block IV. Mapping and bundle adjustment

네 번째 block은 map을 업데이트하고 학습한다.

mapper.process_frame(
    dataset.cur_point_cloud_torch,
    dataset.cur_sem_labels_torch,
    dataset.cur_pose_torch,
    frame_id,
    (config.dynamic_filter_on and frame_id > 0),
)

process_frame()은 크게 네 가지를 한다.

현재 scan에서 SDF training sample을 만든다.
새 observation 근처에 neural point를 추가한다.
training data pool에 현재 sample을 넣는다.
local map을 현재 frame 주변으로 reset한다.

그 다음 실제 optimization을 한다.

mapper.mapping(cur_iter_num)

mapping()은 batch sample을 뽑고, neural point feature와 decoder MLP를 업데이트해서 SDF label을 맞춘다. 초기 frame에서는 더 많이 학습하고, 어느 정도 지나면 decoder를 freeze한다.

if (frame_id-neural_points.reboot_ts) == config.freeze_after_frame:
    freeze_decoders(mlp_dict, config)

내가 보기에는 이 부분이 PIN-SLAM을 이해할 때 아주 중요하다. 초반에는 MLP decoder와 neural point feature가 같이 학습되지만, online SLAM에서 decoder가 계속 바뀌면 과거 map의 의미도 흔들릴 수 있다. 그래서 일정 frame 이후 decoder를 freeze하고, 주로 neural point feature를 통해 map을 업데이트하는 구조로 간다.

Block V. Mesh reconstruction and visualization

visualization이 켜져 있으면 주기적으로 mesh와 SDF slice를 만든다.

cur_mesh = mesher.recon_aabb_collections_mesh(...)
cur_sdf_slice = mesher.generate_bbx_sdf_hor_slice(...)

mesh reconstruction은 SLAM estimation 자체에는 필수는 아니지만, implicit SDF map의 장점을 보여주는 중요한 output이다. point cloud map은 그냥 점들의 집합이지만, SDF field는 공간을 query할 수 있기 때문에 marching cubes로 surface mesh를 뽑을 수 있다.

Block VI. Save results

마지막에는 trajectory, neural point map, mesh를 저장한다.

pose_eval_results = dataset.write_results()
neural_points.prune_map(...)
neural_points.recreate_hash(...)

if config.save_map:
    save_implicit_map(run_path, neural_points, mlp_dict)

if config.save_mesh:
    cur_mesh = mesher.recon_aabb_collections_mesh(...)

여기서 저장되는 implicit map은 단순 .ply가 아니라, neural points와 decoder를 함께 저장해야 한다. 왜냐하면 SDF는 neural_points만으로도, decoder만으로도 정의되지 않고, 둘의 조합으로 정의되기 때문이다.

전체 흐름을 한 줄로 다시 정리

PIN-SLAM의 frame loop는 다음처럼 이해하면 된다.

LiDAR scan
  -> preprocessing
  -> current neural SDF map에 scan-to-map registration
  -> odometry pose update
  -> loop candidate detection + PGO
  -> neural point map deformation after PGO
  -> ray sampling으로 SDF labels 생성
  -> neural point 추가 + local SDF training
  -> mesh / map / trajectory output

이 흐름에서 가장 중요한 두 축은 다음이다.

tracking: map을 고정하고 pose를 찾는다.
mapping: pose를 고정하고 map을 학습한다.

이 둘을 매 frame 번갈아 수행하는 것이 PIN-SLAM의 기본 리듬이다.

처음 읽을 때 헷갈리는 지점

내가 처음 읽으면서 헷갈렸던 지점은 다음과 같다.

NeuralPoints는 MLP가 아니다. Sparse map data structure이고, SDF decode는 Decoder가 한다.
Decoder만 학습하지 않는다. Neural point feature도 nn.Parameter로 같이 optimization된다.
query_feature()는 단순 nearest neighbor가 아니다. Voxel hash, local feature, IDW weight를 한 번에 만든다.
PGO 이후 pose만 update되지 않는다. adjust_map()으로 neural point 위치와 orientation도 같이 움직인다.
Loop detection과 loop registration은 다르다. Descriptor는 후보 탐색, tracker는 transform refinement 담당이다.

다음 글에서는 PIN-SLAM의 핵심 map representation인 NeuralPoints와 Decoder를 읽는다. 특히 update(), reset_local_map(), query_feature()가 어떤 순서로 map을 만들고 query하는지 볼 예정이다.

PIN-SLAM 코드 읽기 시리즈입니다.