3DGS优化强化学习：从零构建高效训练框架的实践指南

1次阅读

共计 2192 个字符，预计需要花费 6 分钟才能阅读完成。

在 3D 场景生成系统（3DGS）中应用强化学习，和传统 RL 任务相比有几个明显的差异点：

状态空间维度爆炸 ：一个中等精度的 3D 场景可能包含上万个点云数据，直接作为状态输入会导致神经网络难以处理
奖励稀疏性 ：比如在室内导航任务中，只有到达目标点才能获得正奖励，中间过程缺乏有效反馈
计算资源消耗大 ：3D 渲染本身就很吃显存，再加上 RL 需要大量交互采样，普通 GPU 根本扛不住

这些特性导致直接用现成的 RL 算法（比如 PPO）训练时，经常会出现训练不稳定、收敛慢的问题。

直接处理原始点云数据效率太低。实践中我们采用 PointNet++ 网络对场景进行特征提取：

class PointNetEncoder(nn.Module):
    def __init__(self, feat_dim=256):
        super().__init__()
        self.sa1 = PointNetSetAbstraction(512, 0.2, 32, 3+3, [64, 64, 128]) 
        self.sa2 = PointNetSetAbstraction(128, 0.4, 32, 128+3, [128, 128, 256])
        self.fc = nn.Linear(256, feat_dim)

    def forward(self, xyz, features):
        # xyz: (B,N,3), features: (B,N,C)
        l1_xyz, l1_features = self.sa1(xyz, features)
        l2_xyz, l2_features = self.sa2(l1_xyz, l1_features)
        return self.fc(l2_features.squeeze(-1))

关键设计点：
– 使用两层 Set Abstraction 逐步降采样
– 最终输出 256 维特征向量，比原始点云小两个数量级
– 保留法向量等几何特征作为输入

针对稀疏奖励问题，我们设计了一个分层奖励系统：

def get_reward(self, state, action):
    # 基础奖励：是否到达目标
    goal_reward = 10.0 if reach_goal else 0.0

    # 进度奖励：朝向目标的移动
    progress = (prev_dist - curr_dist) / max_step
    progress_reward = 2.0 * progress

    # 生存惩罚：鼓励高效探索
    step_penalty = -0.05

    # 碰撞惩罚
    collision_penalty = -1.0 if collision else 0.0

    return goal_reward + progress_reward + step_penalty + collision_penalty

这种设计让智能体在训练早期就能获得有意义的反馈信号。

使用 PyTorch 的 DataParallel 实现多环境并行采样：

class ParallelSampler:
    def __init__(self, env_fn, n_workers=4):
        self.envs = [env_fn() for _ in range(n_workers)]
        self.obs = [env.reset() for env in self.envs]

    def sample(self, policy, steps):
        batch = defaultdict(list)
        for _ in range(steps):
            with torch.no_grad():
                actions = policy(torch.stack(self.obs))

            for i, env in enumerate(self.envs):
                next_obs, reward, done, info = env.step(actions[i])
                batch['obs'].append(self.obs[i])
                batch['actions'].append(actions[i])
                # ... 其他数据收集

                self.obs[i] = env.reset() if done else next_obs

        return {k: torch.stack(v) for k,v in batch.items()}

在 NVIDIA RTX 3090 上的测试结果：