3D 视觉与空间理解：从 NeRF 到世界模型

import torch
import torch.nn as nn
import torch.nn.functional as F

class NeRF(nn.Module):
    """NeRF 核心网络：位置编码 + MLP"""
    
    def __init__(self, hidden_dim=256, num_layers=8, 
                 pos_encoding_freq=10, view_encoding_freq=4):
        super().__init__()
        self.pos_encoding_freq = pos_encoding_freq
        self.view_encoding_freq = view_encoding_freq
        
        # 位置编码维度: 3 * (2 * freq + 1)
        pos_dim = 3 * (2 * pos_encoding_freq + 1)
        view_dim = 3 * (2 * view_encoding_freq + 1)
        
        # 主干网络：位置 -> 密度 + 特征
        layers = []
        layers.append(nn.Linear(pos_dim, hidden_dim))
        for _ in range(num_layers - 1):
            layers.append(nn.ReLU())
            layers.append(nn.Linear(hidden_dim, hidden_dim))
        self.feature_layers = nn.Sequential(*layers)
        
        # 密度输出
        self.density_head = nn.Linear(hidden_dim, 1)
        
        # 颜色网络：特征 + 视角 -> RGB
        self.color_head = nn.Sequential(
            nn.Linear(hidden_dim + view_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 3),
            nn.Sigmoid()
        )
    
    def positional_encoding(self, x, freq):
        """正弦位置编码"""
        enc = [x]
        for i in range(freq):
            enc.append(torch.sin(x * (2 ** i)))
            enc.append(torch.cos(x * (2 ** i)))
        return torch.cat(enc, dim=-1)
    
    def forward(self, points, view_dirs):
        """
        points: [N, 3] 3D 空间坐标
        view_dirs: [N, 3] 观察方向
        返回: rgb [N, 3], density [N, 1]
        """
        pos_enc = self.positional_encoding(points, self.pos_encoding_freq)
        view_enc = self.positional_encoding(view_dirs, self.view_encoding_freq)
        
        features = self.feature_layers(pos_enc)
        density = self.density_head(features)
        color = self.color_head(torch.cat([features, view_enc], dim=-1))
        
        return color, density

# 体积渲染：沿射线积分
def volume_render(colors, densities, deltas, dists):
    """简化的体积渲染"""
    alphas = 1 - torch.exp(-densities * deltas)
    weights = alphas * torch.cumprod(
        torch.cat([torch.ones_like(alphas[:1]), 1 - alphas[:-1]], dim=0)
    )
    rgb = torch.sum(weights * colors, dim=0)
    return rgb, weights

# NeRF 训练循环示例
def train_nerf(model, rays, gt_colors, optimizer, num_rays=4096):
    """单步训练"""
    optimizer.zero_grad()
    
    # 随机采样射线
    indices = torch.randperm(len(rays))[:num_rays]
    sampled_rays = rays[indices]
    sampled_gt = gt_colors[indices]
    
    # 沿射线采样点
    points, view_dirs, deltas = sample_points_on_rays(sampled_rays)
    
    # 前向传播
    pred_colors, densities = model(points, view_dirs)
    
    # 体积渲染
    rendered_colors, _ = volume_render(
        pred_colors, densities, deltas, 
        dists=sampled_rays.near_far_dist
    )
    
    # MSE 损失
    loss = F.mse_loss(rendered_colors, sampled_gt)
    loss.backward()
    optimizer.step()
    
    return loss.item()

import torch
from dataclasses import dataclass

@dataclass
class Gaussian3D:
    """单个 3D 高斯"""
    mean: torch.Tensor          # [3] 位置
    covariance: torch.Tensor    # [3, 3] 协方差矩阵
    opacity: torch.Tensor       # [1] 不透明度
    sh_coeffs: torch.Tensor     # [N, 3] 球谐系数（颜色）

class GaussianSplatting:
    """3D 高斯溅射渲染器"""
    
    def __init__(self, num_gaussians: int, sh_degree: int = 3):
        self.num_gaussians = num_gaussians
        self.sh_degree = sh_degree
        self.num_sh_coeffs = (sh_degree + 1) ** 2
        
        # 初始化高斯参数
        self.means = nn.Parameter(torch.randn(num_gaussians, 3) * 0.1)
        # 使用旋转（四元数）和缩放参数化协方差
        self.rotations = nn.Parameter(torch.randn(num_gaussians, 4))
        self.scales = nn.Parameter(torch.zeros(num_gaussians, 3))
        self.opacities = nn.Parameter(torch.zeros(num_gaussians, 1))
        self.sh_coeffs = nn.Parameter(torch.zeros(
            num_gaussians, self.num_sh_coeffs, 3
        ))
    
    def get_covariance(self) -> torch.Tensor:
        """从旋转和缩放构建协方差矩阵"""
        # 归一化四元数
        q = F.normalize(self.rotations, dim=-1)
        # 构建旋转矩阵
        R = quaternion_to_matrix(q)
        # 缩放矩阵
        S = torch.diag_embed(torch.exp(self.scales))
        # 协方差 = R * S^2 * R^T
        return R @ (S @ S.transpose(-1, -2)) @ R.transpose(-1, -2)
    
    def project_to_camera(self, means, view_matrix, focal, center):
        """将 3D 高斯投影到相机坐标系"""
        # 转换到相机坐标
        means_cam = (view_matrix[:3, :3] @ means.T).T + view_matrix[:3, 3]
        # 透视投影
        depth = means_cam[:, 2:3]
        means_proj = means_cam[:, :2] / depth * focal + center
        # 协方差投影（使用雅可比近似）
        return means_proj, depth

    def render(self, camera_params) -> torch.Tensor:
        """完整渲染管线"""
        means_proj, depth = self.project_to_camera(
            self.means, camera_params.view_matrix,
            camera_params.focal, camera_params.center
        )
        cov_2d = project_covariance(self.get_covariance(), camera_params)
        # 深度排序 + alpha 混合
        return tile_based_rasterization(
            means_proj, cov_2d, torch.sigmoid(self.opacities),
            self.sh_coeffs, depth
        )

from dataclasses import dataclass
from typing import List, Optional

@dataclass
class SpatialUnderstanding:
    """3D 场景的空间理解结果"""
    objects: List[dict]        # 检测到的物体及其 3D 位姿
    surfaces: List[dict]       # 平面（地面、墙面、桌面）
    relationships: List[dict]  # 空间关系（在...上面、在...旁边）
    free_space: List[dict]     # 可行走/可操作区域
    scene_graph: dict          # 场景图结构

class SpatialLLM:
    """3D 空间理解大语言模型"""
    
    def __init__(self, base_llm, point_encoder, spatial_head):
        self.llm = base_llm
        self.point_encoder = point_encoder  # 点云 -> token
        self.spatial_head = spatial_head    # 空间推理头
    
    def understand_scene(self, point_cloud: torch.Tensor,
                         instruction: str) -> SpatialUnderstanding:
        """理解 3D 场景并回答空间问题"""
        # 编码点云为 token
        point_tokens = self.point_encoder(point_cloud)
        
        # 拼接指令 token 和点云 token
        text_tokens = self.llm.tokenize(instruction)
        combined_tokens = torch.cat([text_tokens, point_tokens], dim=1)
        
        # 推理
        output = self.llm(combined_tokens)
        
        # 解析空间理解结果
        return self.spatial_head.parse(output)
    
    def plan_manipulation(self, scene: SpatialUnderstanding,
                          task: str) -> List[dict]:
        """规划机器人操作动作序列"""
        prompt = f"""
        场景理解: {scene}
        任务: {task}
        
        请生成操作动作序列，包括：
        1. 移动路径
        2. 抓取目标
        3. 放置位置
        4. 碰撞避免
        """
        return self.llm.generate_action_sequence(prompt)