`depth_fm.objectives.losses`¶

losses ¶

Loss functions for Mars DepthFM training.

[REFACTOR NOTE] 1226 LOC; each loss class is 200+ LOC with internal helpers. Candidate for future split into one file per loss (e.g. photoclinometric.py, absolute_depth.py, laplacian.py, ordinal_ranking.py, combined.py). Out of scope for the structural refactor that produced this layout.

Conditional Flow Matching velocity regression

L_FM = || v_theta(z_t, t; z_img) - (z_depth - z_img) ||^2

Auxiliary losses (all operate in pixel space on the predicted clean depth): Surface normals consistency — angular error between Sobel normals Focal Frequency Loss (FFL) — penalises frequency-domain discrepancies Multi-Scale Gradient Loss — L1 slope error at 1×, 2×, 4× downsampling

PhotoclinometricLoss ¶

PhotoclinometricLoss(ssim_weight: float = 0.5, scales: tuple[int, ...] = (1, 2, 4))

Bases: Module

State-of-the-art photoclinometric loss for planetary DTM estimation.

Improvements over simple Lambertian

Lunar-Lambert reflectance model (McEwen 1991; Hapke 2012): Blends Lambertian and Lommel-Seeliger components with a learnable weight, correctly modeling the limb-darkening behaviour of regolith.
Multi-scale rendering at 1×, 2×, 4× — enforces macro-scale topographic consistency alongside fine detail.
Combined Pearson + z-normalised SSIM loss. Pearson captures global correlation; SSIM captures local structural similarity (Wang et al. 2004). Both components are applied to per-image z-score normalised inputs, making the loss invariant to global luminance and contrast mismatch between the render and the ortho. This matters here because the scene intensity/ambient parameters are fitted with a pure Lambertian OLS, while the render uses Lunar-Lambert — so a systematic luminance offset is expected and should not be penalised.
Unclamped render. Because both loss terms are scale- and shift-invariant, clamping the render to [-1, 1] would only destroy gradient information in saturated shadow pixels — which is exactly where photoclinometric information is densest.
Variance floor and proper epsilon handling to prevent NaN gradients when renders are near-flat (e.g. shadow regions).

The loss is

L = (1 - α) * (1 - Pearson) + α * (1 - SSIM_z)

averaged over scales, where α = 0.5 by default and SSIM_z is SSIM computed on z-score normalised inputs.

NaN-safety

Every code path is hardened against NaN under DDP + gradient accumulation + mixed precision. Key invariants: - NEVER early-return with a detached zero — always flow through self.lunar_lambert_logit so DDP gradient sync sees every param. - All internal math forced to float32 to avoid bf16/fp16 overflow from the spatial_scale multiplication (up to 256×). - sqrt() always gets eps inside to prevent ∞ gradients. - F.normalize always gets eps > 0. - Lommel-Seeliger cos_e clamped to prevent denominator collapse.

Source code in src/depth_fm/objectives/losses.py

def __init__(self, ssim_weight: float = 0.5, scales: tuple[int, ...] = (1, 2, 4)):
    super().__init__()
    self.ssim_weight = ssim_weight
    self.scales = scales

    # Sobel kernels for computing surface gradients
    sobel_x = torch.tensor([[-1., 0., 1.],
                            [-2., 0., 2.],
                            [-1., 0., 1.]]) / 8.0
    sobel_y = torch.tensor([[-1., -2., -1.],
                            [0., 0., 0.],
                            [1., 2., 1.]]) / 8.0

    self.register_buffer("kernel_x", sobel_x.view(1, 1, 3, 3))
    self.register_buffer("kernel_y", sobel_y.view(1, 1, 3, 3))

    # Learnable Lunar-Lambert blend: sigmoid(raw) → L ∈ (0, 1)
    # L=1 → pure Lambertian, L=0 → pure Lommel-Seeliger
    # Initialised at 0.0 → sigmoid(0)=0.5 (equal blend)
    self.lunar_lambert_logit = nn.Parameter(torch.tensor(0.0))

surface_normals ¶

surface_normals(depth: Tensor) -> torch.Tensor

Public: compute unit surface normals from a depth map.

Source code in src/depth_fm/objectives/losses.py

def surface_normals(self, depth: torch.Tensor) -> torch.Tensor:
    """Public: compute unit surface normals from a depth map."""
    return self._get_surface_normals(depth.float())

render_from_depth ¶

render_from_depth(depth: Tensor, sun_vector: Tensor, intensity: Tensor, ambient: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Public: render a depth map under the current Lunar-Lambert model.

Handles input normalisation/broadcasting so callers can pass per-sample sun vectors and scalar exposure parameters directly.

Returns:

Name	Type	Description
`render`	`Tensor`	(B, 1, H, W) unclamped radiance.
`normals`	`Tensor`	(B, 3, H, W) unit surface normals.

Source code in src/depth_fm/objectives/losses.py

def render_from_depth(
        self, depth: torch.Tensor, sun_vector: torch.Tensor,
        intensity: torch.Tensor, ambient: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Public: render a depth map under the current Lunar-Lambert model.

    Handles input normalisation/broadcasting so callers can pass
    per-sample sun vectors and scalar exposure parameters directly.

    Returns:
        render:  (B, 1, H, W) unclamped radiance.
        normals: (B, 3, H, W) unit surface normals.
    """
    B = depth.shape[0]
    depth = depth.float()
    normals = self._get_surface_normals(depth)
    l_dir = F.normalize(sun_vector.float(), p=2, dim=1, eps=1e-6).view(B, 3, 1, 1)
    intensity_b = intensity.float().reshape(B, 1, 1, 1)
    ambient_b = ambient.float().reshape(B, 1, 1, 1)
    render = self._lunar_lambert_render(normals, l_dir, intensity_b, ambient_b)
    return render, normals

FlowMatchingVelocityLoss ¶

FlowMatchingVelocityLoss(use_confidence_weighting: bool = False)

Bases: Module

Conditional flow matching velocity loss.

Given a predicted velocity v_pred and the target velocity (z_depth - z_img), compute the L2 loss, optionally weighted by confidence maps.

Source code in src/depth_fm/objectives/losses.py

def __init__(self, use_confidence_weighting: bool = False):
    super().__init__()
    self.use_confidence = use_confidence_weighting

forward ¶

forward(v_pred: Tensor, v_target: Tensor, confidence: Tensor = None) -> torch.Tensor

Parameters:

Name	Type	Description	Default
`v_pred`	`Tensor`	predicted velocity (B, C, h, w)	required
`v_target`	`Tensor`	target velocity z_depth - z_img (B, C, h, w)	required
`confidence`	`Tensor`	optional weight map (B, 1, H, W) — will be downsampled	`None`

Returns:

Type	Description
`Tensor`	scalar loss

Source code in src/depth_fm/objectives/losses.py

def forward(
        self,
        v_pred: torch.Tensor,
        v_target: torch.Tensor,
        confidence: torch.Tensor = None,
) -> torch.Tensor:
    """
    Args:
        v_pred: predicted velocity (B, C, h, w)
        v_target: target velocity z_depth - z_img (B, C, h, w)
        confidence: optional weight map (B, 1, H, W) — will be downsampled

    Returns:
        scalar loss
    """
    sq_error = (v_pred.float() - v_target.float()).pow(2)

    if self.use_confidence and confidence is not None:
        # Area interpolation gives a fractional-valid weight per latent
        # cell: if 3 of 4 original pixels in the receptive field are
        # valid, the weight is 0.75. This is more faithful than nearest
        # neighbor because a single valid pixel shouldn't cause the
        # whole latent cell to count as fully valid.
        mask = F.interpolate(
            confidence.float(),
            size=v_pred.shape[-2:],
            mode="area",
        )
        sq_error = sq_error * mask

        # Normalize strictly by the active area to maintain gradient scale
        active_elements = mask.sum() * v_pred.shape[1]
        return sq_error.sum() / (active_elements + 1e-8)

    return sq_error.mean()

SurfaceNormalsLoss ¶

SurfaceNormalsLoss()

Bases: Module

Surface normals consistency loss.

Computes surface normals from the spatial gradients (Sobel) of predicted and GT depth maps, then measures angular difference via cosine similarity.

This operates in PIXEL space (decoded depth), not latent space.

DepthFM uses the "normal consistency" formulation from: Fan et al., "Three-filters-to-normal", IEEE RA-L 2021 which computes normals using finite differences and compares with cosine loss.

Source code in src/depth_fm/objectives/losses.py

def __init__(self):
    super().__init__()
    sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32)
    sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32)
    self.register_buffer("sobel_x", sobel_x.view(1, 1, 3, 3))
    self.register_buffer("sobel_y", sobel_y.view(1, 1, 3, 3))

MultiScaleGradientLoss ¶

MultiScaleGradientLoss(scales: tuple[int, ...] = (1, 2, 4))

Bases: Module

Multi-scale gradient matching loss for high-frequency topography detail.

Penalises L1 differences between predicted and GT spatial gradients at multiple downsampling scales. This is inherently scale-invariant (derivative differences, not absolute values) and forces the model to learn crisp slopes, crater rims, and dune ripples regardless of residual absolute-scale errors.

Loss

L_grad = mean over scales of: mean(|∇x_pred - ∇x_gt| + |∇y_pred - ∇y_gt|)

Parameters:

Name	Type	Description	Default
`scales`	`tuple[int, ...]`	tuple of integer downsampling factors (applied as average pooling).	`(1, 2, 4)`

Source code in src/depth_fm/objectives/losses.py

def __init__(self, scales: tuple[int, ...] = (1, 2, 4)):
    super().__init__()
    self.scales = scales

FocalFrequencyLoss ¶

FocalFrequencyLoss(loss_weight: float = 1.0, alpha: float = 1.0)

Bases: Module

Wrapper around the focal-frequency-loss package.

Penalises discrepancies in the Fourier domain, explicitly forcing the network to recover high-frequency details that MSE smooths over.

Requires: pip install focal-frequency-loss Paper: https://arxiv.org/pdf/2012.12821

Parameters:

Name	Type	Description	Default
`loss_weight`	`float`	Overall weight for this loss term.	`1.0`
`alpha`	`float`	Focal factor — higher values focus more on hard frequencies.	`1.0`

Source code in src/depth_fm/objectives/losses.py

def __init__(self, loss_weight: float = 1.0, alpha: float = 1.0):
    super().__init__()
    if not _HAS_FFL:
        raise ImportError(
            "focal-frequency-loss is not installed. "
            "Run: pip install focal-frequency-loss"
        )
    self._ffl = _FFL(loss_weight=loss_weight, alpha=alpha)

AbsoluteDepthLoss ¶

AbsoluteDepthLoss(delta: float = 0.1)

Bases: Module

Direct Huber regression in [-1, 1] depth space.

Fills the 'absolute accuracy' gap: every other pixel-space loss in the pipeline is scale/shift-invariant (normals, gradient, photoclinometric). Without a direct loss on absolute values, a model could have perfect slopes and renders but systematically wrong elevations.

Huber combines

L2 behaviour near convergence (|r| <= delta): smooth gradients, fast final-stage descent.
L1 behaviour in the tails (|r| > delta): robust to outliers from stereo GT artefacts and nodata boundaries.

For inputs bounded in [-1, 1], delta=0.1 treats errors under ~10 % of the data range as 'near optimum' (L2) and larger errors as 'outliers' (L1).

Source code in src/depth_fm/objectives/losses.py

def __init__(self, delta: float = 0.1):
    super().__init__()
    self.delta = delta

per_pixel_loss ¶

per_pixel_loss(pred: Tensor, gt: Tensor) -> torch.Tensor

Un-reduced Huber loss map, shape (B, 1, H, W).

Source code in src/depth_fm/objectives/losses.py

def per_pixel_loss(self, pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
    """Un-reduced Huber loss map, shape (B, 1, H, W)."""
    if pred.shape[1] == 3:
        pred = pred[:, :1]
    if gt.shape[1] == 3:
        gt = gt[:, :1]
    return F.huber_loss(
        pred.float(), gt.float(), reduction="none", delta=self.delta
    )

LaplacianLoss ¶

LaplacianLoss()

Bases: Module

Second-order curvature consistency via discrete Laplacian.

Captures shape information that first-order gradients miss. Critical for Mars because its defining landforms are curvature features: - Crater bowls (negative curvature) - Central peaks (positive curvature) - Volcanic calderas, rims, scarps

A model matching all first-order gradients could still get crater concavity qualitatively wrong; the Laplacian forces correct concave/convex structure.

L = mean_over_valid(|∇²pred - ∇²gt|)

Source code in src/depth_fm/objectives/losses.py

def __init__(self):
    super().__init__()
    # Discrete 5-point Laplacian
    kernel = torch.tensor(
        [[0.0, 1.0, 0.0], [1.0, -4.0, 1.0], [0.0, 1.0, 0.0]]
    ).view(1, 1, 3, 3)
    self.register_buffer("kernel", kernel)

laplacian ¶

laplacian(depth: Tensor) -> torch.Tensor

Discrete Laplacian (∇²d) of a depth map, shape (B, 1, H, W).

Source code in src/depth_fm/objectives/losses.py

def laplacian(self, depth: torch.Tensor) -> torch.Tensor:
    """Discrete Laplacian (∇²d) of a depth map, shape (B, 1, H, W)."""
    if depth.shape[1] == 3:
        depth = depth[:, :1]
    padded = F.pad(depth.float(), (1, 1, 1, 1), mode="replicate")
    return F.conv2d(padded, self.kernel.float())

per_pixel_loss ¶

per_pixel_loss(pred: Tensor, gt: Tensor) -> torch.Tensor

Absolute Laplacian difference, shape (B, 1, H, W).

Source code in src/depth_fm/objectives/losses.py

def per_pixel_loss(self, pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
    """Absolute Laplacian difference, shape (B, 1, H, W)."""
    return (self.laplacian(pred) - self.laplacian(gt)).abs()

OrdinalRankingLoss ¶

OrdinalRankingLoss(margin: float = 0.02, num_pairs: int = 10000)

Bases: Module

Relative-ordering consistency over sampled pixel pairs.

For each sampled pair (i, j): if gt[i] - gt[j] > margin then pred[i] should exceed pred[j]. Uses a hinge loss:

L_pair = ReLU( -sign(gt_i - gt_j) * (pred_i - pred_j) )

averaged over pairs where |gt_i - gt_j| > margin (i.e. pairs whose GT ordering is unambiguous).

Why this helps on Mars¶

Mars DTM ground truth from stereo reconstruction has systematic noise (registration, interpolation, crater-wall occlusion). Absolute values in noisy regions are unreliable, but pairwise ORDERING is robust — if crater floor is lower than rim in the GT, it is almost certainly lower in reality even if the absolute depths are off. Ordinal loss gives a supervision signal that degrades gracefully with GT noise.

Based on: Chen et al., 'Single-Image Depth Perception in the Wild', NeurIPS 2016 (DIW) — adapted here to dense GT by random pair sampling.

Source code in src/depth_fm/objectives/losses.py

def __init__(self, margin: float = 0.02, num_pairs: int = 10000):
    super().__init__()
    self.margin = margin
    self.num_pairs = num_pairs

sample_pairs ¶

sample_pairs(B: int, N: int, device: device, generator: Generator | None = None, confidence: Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]

Sample (idx_i, idx_j) flat indices, each shape (B, num_pairs).

Deterministic if a generator is provided — used by viz for reproducible figures.

Source code in src/depth_fm/objectives/losses.py

def sample_pairs(
        self,
        B: int,
        N: int,
        device: torch.device,
        generator: torch.Generator | None = None,
        confidence: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Sample (idx_i, idx_j) flat indices, each shape (B, num_pairs).

    Deterministic if a generator is provided — used by viz for
    reproducible figures.
    """
    if confidence is not None:
        # Flatten to (B, N) and ensure non-negative weights
        weights = torch.clamp(confidence.reshape(B, N).float(), min=0.0)

        # Add a tiny epsilon. If a batch element has completely zero confidence,
        # this prevents a multinomial crash by falling back to uniform sampling.
        # The invalid pairs will still be safely ignored in pair_stats.
        weights = weights + 1e-8

        idx = torch.multinomial(
            weights,
            self.num_pairs * 2,
            replacement=True,
            generator=generator
        )
    else:
        if generator is not None:
            idx = torch.randint(
                N, (B, self.num_pairs * 2), device=device, generator=generator
            )
        else:
            idx = torch.randint(N, (B, self.num_pairs * 2), device=device)

    return idx[:, : self.num_pairs], idx[:, self.num_pairs:]

pair_stats ¶

pair_stats(pred: Tensor, gt: Tensor, idx_i: Tensor, idx_j: Tensor, confidence: Tensor = None) -> dict

Per-pair losses, ordering labels, violation flags.

Returns dict with

losses (B, num_pairs) — hinge loss, 0 for ambiguous pairs ordered (B, num_pairs) bool — |gt_diff| > margin and valid violations (B, num_pairs) bool — ordered and pred disagrees gt_diff (B, num_pairs) — gt_i - gt_j pr_diff (B, num_pairs) — pred_i - pred_j

Source code in src/depth_fm/objectives/losses.py

def pair_stats(
        self,
        pred: torch.Tensor,
        gt: torch.Tensor,
        idx_i: torch.Tensor,
        idx_j: torch.Tensor,
        confidence: torch.Tensor = None,
) -> dict:
    """Per-pair losses, ordering labels, violation flags.

    Returns dict with:
      losses      (B, num_pairs) — hinge loss, 0 for ambiguous pairs
      ordered     (B, num_pairs) bool — |gt_diff| > margin and valid
      violations  (B, num_pairs) bool — ordered and pred disagrees
      gt_diff     (B, num_pairs) — gt_i - gt_j
      pr_diff     (B, num_pairs) — pred_i - pred_j
    """
    if pred.shape[1] == 3:
        pred = pred[:, :1]
    if gt.shape[1] == 3:
        gt = gt[:, :1]

    B = pred.shape[0]
    pred_flat = pred.reshape(B, -1).float()
    gt_flat = gt.reshape(B, -1).float()

    gt_i = gt_flat.gather(1, idx_i)
    gt_j = gt_flat.gather(1, idx_j)
    pr_i = pred_flat.gather(1, idx_i)
    pr_j = pred_flat.gather(1, idx_j)

    gt_diff = gt_i - gt_j
    pr_diff = pr_i - pr_j
    sign_gt = torch.sign(gt_diff)

    ordered = gt_diff.abs() > self.margin

    # We keep this check even with multinomial sampling to catch
    # the 1e-8 epsilon fallback edge-case where a batch is totally unconfident.
    if confidence is not None:
        conf_flat = confidence.reshape(B, -1).float()
        c_i = conf_flat.gather(1, idx_i)
        c_j = conf_flat.gather(1, idx_j)
        ordered = ordered & (c_i > 0.5) & (c_j > 0.5)

    losses = F.relu(-sign_gt * pr_diff) * ordered.float()
    # 'violation' = ordered AND predicted difference has wrong sign
    violations = ordered & (sign_gt * pr_diff <= 0)

    return {
        "losses": losses,
        "ordered": ordered,
        "violations": violations,
        "gt_diff": gt_diff,
        "pr_diff": pr_diff,
    }

CombinedLoss ¶

CombinedLoss(velocity_weight: float = 1.0, normals_weight: float = 0.1, normals_start_step: int = 2000, use_confidence_weighting: bool = False, freq_weight: float = 0.0, freq_start_step: int = 0, freq_alpha: float = 1.0, grad_weight: float = 0.0, grad_start_step: int = 0, grad_scales: tuple[int, ...] = (1, 2, 4), photo_weight: float = 0.0, photo_start_step: int = 2000, huber_weight: float = 3.0, huber_start_step: int = 0, huber_delta: float = 0.1, laplacian_weight: float = 0.05, laplacian_start_step: int = 0, ordinal_weight: float = 1.0, ordinal_start_step: int = 0, ordinal_margin: float = 0.02, ordinal_num_pairs: int = 10000)

Bases: Module

Combined training loss for Mars DepthFM.

L_total = w_vel * L_FM + w_norm * L_normals (after normals_start_step) + w_freq * L_FFL (after freq_start_step) + w_grad * L_grad (after grad_start_step) + w_photo * L_photo (after photo_start_step) + w_huber * L_Huber (after huber_start_step) [NEW] + w_lap * L_Laplacian (after laplacian_start_step) [NEW] + w_ord * L_Ordinal (after ordinal_start_step) [NEW]

Pixel-space losses (normals, FFL, grad, photo, huber, laplacian, ordinal) require the caller to pass pred_depth_pixels and gt_depth_pixels decoded WITHOUT no_grad so that gradients flow back through the decoder to v_pred.

Source code in src/depth_fm/objectives/losses.py

def __init__(
        self,
        velocity_weight: float = 1.0,
        normals_weight: float = 0.1,
        normals_start_step: int = 2000,
        use_confidence_weighting: bool = False,
        freq_weight: float = 0.0,
        freq_start_step: int = 0,
        freq_alpha: float = 1.0,
        grad_weight: float = 0.0,
        grad_start_step: int = 0,
        grad_scales: tuple[int, ...] = (1, 2, 4),
        photo_weight: float = 0.0,
        photo_start_step: int = 2000,
        huber_weight: float = 3.0,
        huber_start_step: int = 0,
        huber_delta: float = 0.1,
        laplacian_weight: float = 0.05,
        laplacian_start_step: int = 0,
        ordinal_weight: float = 1.0,
        ordinal_start_step: int = 0,
        ordinal_margin: float = 0.02,
        ordinal_num_pairs: int = 10000,
):
    super().__init__()
    # ---- Photo / velocity / existing aux ----------------------------
    self.photo_start_step = photo_start_step
    self.photo_weight = photo_weight
    self.photo_loss = PhotoclinometricLoss(
        ssim_weight=0.5,
        scales=(1, 2, 4),
    )
    self.velocity_loss = FlowMatchingVelocityLoss(use_confidence_weighting)
    self.normals_loss = SurfaceNormalsLoss()
    self.grad_loss = MultiScaleGradientLoss(scales=grad_scales) if grad_weight > 0 else None
    self.vel_weight = velocity_weight
    self.norm_weight = normals_weight
    self.norm_start = normals_start_step
    self.freq_weight = freq_weight
    self.freq_start = freq_start_step
    self.grad_weight = grad_weight
    self.grad_start = grad_start_step
    self.use_confidence_weighting = use_confidence_weighting

    # ---- New losses --------------------------------------------------
    self.huber_weight = huber_weight
    self.huber_start = huber_start_step
    self.huber_loss = AbsoluteDepthLoss(delta=huber_delta) if huber_weight > 0 else None

    self.laplacian_weight = laplacian_weight
    self.laplacian_start = laplacian_start_step
    self.laplacian_loss = LaplacianLoss() if laplacian_weight > 0 else None

    self.ordinal_weight = ordinal_weight
    self.ordinal_start = ordinal_start_step
    self.ordinal_loss = (
        OrdinalRankingLoss(margin=ordinal_margin, num_pairs=ordinal_num_pairs)
        if ordinal_weight > 0 else None
    )

    if self.use_confidence_weighting:
        logging.info("Using confidence weighting to improve nodata region filtering")

    if freq_weight > 0:
        if not _HAS_FFL:
            logger.warning(
                "freq_weight=%.2f requested but focal-frequency-loss is not installed. "
                "FFL will be skipped. Run: pip install focal-frequency-loss",
                freq_weight,
            )
            self.ffl = None
        else:
            self.ffl = FocalFrequencyLoss(loss_weight=1.0, alpha=freq_alpha)
    else:
        self.ffl = None

needs_pixel_decode ¶

needs_pixel_decode(global_step: int) -> bool

Return True if any pixel-space loss is active at this step.

Source code in src/depth_fm/objectives/losses.py

def needs_pixel_decode(self, global_step: int) -> bool:
    """Return True if any pixel-space loss is active at this step."""
    return (
            (self.norm_weight > 0 and global_step >= self.norm_start) or
            (self.freq_weight > 0 and self.ffl is not None and global_step >= self.freq_start) or
            (self.grad_weight > 0 and self.grad_loss is not None and global_step >= self.grad_start) or
            (self.photo_weight > 0 and self.photo_loss is not None and global_step >= self.photo_start_step) or
            (self.huber_weight > 0 and self.huber_loss is not None and global_step >= self.huber_start) or
            (
                    self.laplacian_weight > 0 and self.laplacian_loss is not None and global_step >= self.laplacian_start) or
            (self.ordinal_weight > 0 and self.ordinal_loss is not None and global_step >= self.ordinal_start)
    )

forward ¶

forward(v_pred: Tensor, v_target: Tensor, pred_depth_pixels: Tensor = None, gt_depth_pixels: Tensor = None, pred_depth_physical: Tensor = None, confidence: Tensor = None, real_ortho: Tensor | None = None, sun_vector: Tensor | None = None, global_step: int = 0, ambient: Tensor = None, intensity: Tensor = None) -> dict

Parameters:

Name	Type	Description	Default
`v_pred`	`Tensor`	predicted velocity in latent space (B, C, h, w)	required
`v_target`	`Tensor`	target velocity (B, C, h, w)	required
`pred_depth_pixels`	`Tensor`	decoded predicted clean depth (B, 3, H, W)	`None`
`gt_depth_pixels`	`Tensor`	decoded GT depth (B, 3, H, W)	`None`
`confidence`	`Tensor`	optional confidence map (B, 1, H, W)	`None`
`global_step`	`int`	current training step	`0`

Returns:

Type	Description
`dict`	dict with component losses and "total"

Source code in src/depth_fm/objectives/losses.py

def forward(
        self,
        v_pred: torch.Tensor,
        v_target: torch.Tensor,
        pred_depth_pixels: torch.Tensor = None,
        gt_depth_pixels: torch.Tensor = None,
        pred_depth_physical: torch.Tensor = None,
        confidence: torch.Tensor = None,
        real_ortho: torch.Tensor | None = None,
        sun_vector: torch.Tensor | None = None,
        global_step: int = 0,
        ambient: torch.Tensor = None,
        intensity: torch.Tensor = None,
) -> dict:
    """
    Args:
        v_pred: predicted velocity in latent space (B, C, h, w)
        v_target: target velocity (B, C, h, w)
        pred_depth_pixels: decoded predicted clean depth (B, 3, H, W)
        gt_depth_pixels:   decoded GT depth (B, 3, H, W)
        confidence: optional confidence map (B, 1, H, W)
        global_step: current training step

    Returns:
        dict with component losses and "total"
    """
    active_conf = confidence if self.use_confidence_weighting else None

    l_vel = self.velocity_loss(v_pred, v_target, active_conf)

    loss_dict = {
        "velocity": l_vel,
        "normals": torch.tensor(0.0, device=l_vel.device),
        "freq": torch.tensor(0.0, device=l_vel.device),
        "grad": torch.tensor(0.0, device=l_vel.device),
        "photo": torch.tensor(0.0, device=v_pred.device),
        "huber": torch.tensor(0.0, device=v_pred.device),
        "laplacian": torch.tensor(0.0, device=v_pred.device),
        "ordinal": torch.tensor(0.0, device=v_pred.device),
    }

    total = self.vel_weight * l_vel
    has_pixels = pred_depth_pixels is not None and gt_depth_pixels is not None

    # ---- Existing pixel-space losses --------------------------------
    if self.norm_weight > 0 and global_step >= self.norm_start and has_pixels:
        l_norm = self.normals_loss(pred_depth_pixels, gt_depth_pixels, active_conf)
        loss_dict["normals"] = l_norm
        total = total + self.norm_weight * l_norm

    if (
            self.freq_weight > 0
            and self.ffl is not None
            and global_step >= self.freq_start
            and has_pixels
    ):
        l_freq = self.ffl(pred_depth_pixels, gt_depth_pixels, active_conf)
        loss_dict["freq"] = l_freq
        total = total + self.freq_weight * l_freq

    if (
            self.grad_weight > 0
            and self.grad_loss is not None
            and global_step >= self.grad_start
            and has_pixels
    ):
        l_grad = self.grad_loss(pred_depth_pixels, gt_depth_pixels, active_conf)
        loss_dict["grad"] = l_grad
        total = total + self.grad_weight * l_grad

    if (
            self.photo_weight > 0
            and self.photo_loss is not None
            and global_step >= self.photo_start_step
            and has_pixels
    ):
        if sun_vector is None:
            sun_vector = torch.tensor([[0.5, -0.5, 1.0]], device=pred_depth_pixels.device)
            sun_vector = sun_vector.expand(pred_depth_pixels.shape[0], -1)

        B_photo = pred_depth_pixels.shape[0]
        if ambient is None:
            ambient = torch.zeros(B_photo, 1, device=pred_depth_pixels.device)
        if intensity is None:
            intensity = torch.ones(B_photo, 1, device=pred_depth_pixels.device)

        pred_for_photo = pred_depth_physical if pred_depth_physical is not None else pred_depth_pixels
        l_photo = self.photo_loss(
            pred_depth=pred_for_photo,
            real_ortho=real_ortho,
            # Use active_conf (gated by use_confidence_weighting) for
            # consistency with all other aux losses. If the switch is
            # off, fall back to all-ones so photo loss still runs.
            mask=active_conf if active_conf is not None
            else torch.ones_like(pred_depth_pixels[:, :1]),
            sun_vectors=sun_vector,
            ambient=ambient,
            intensity=intensity,
        )

        loss_dict["photo"] = l_photo
        total = total + self.photo_weight * l_photo

    # ---- New pixel-space losses -------------------------------------
    if (
            self.huber_weight > 0
            and self.huber_loss is not None
            and global_step >= self.huber_start
            and has_pixels
    ):
        l_huber = self.huber_loss(pred_depth_pixels, gt_depth_pixels, active_conf)
        loss_dict["huber"] = l_huber
        total = total + self.huber_weight * l_huber

    if (
            self.laplacian_weight > 0
            and self.laplacian_loss is not None
            and global_step >= self.laplacian_start
            and has_pixels
    ):
        l_lap = self.laplacian_loss(pred_depth_pixels, gt_depth_pixels, active_conf)
        loss_dict["laplacian"] = l_lap
        total = total + self.laplacian_weight * l_lap

    if (
            self.ordinal_weight > 0
            and self.ordinal_loss is not None
            and global_step >= self.ordinal_start
            and has_pixels
    ):
        l_ord = self.ordinal_loss(pred_depth_pixels, gt_depth_pixels, active_conf)
        loss_dict["ordinal"] = l_ord
        total = total + self.ordinal_weight * l_ord

    loss_dict["total"] = total

    # ---- Final NaN guard --------------------------------------------
    if not torch.isfinite(total):
        logger.warning(
            "CombinedLoss: non-finite total detected at step %d. "
            "Components: vel=%.4f norm=%.4f freq=%.4f grad=%.4f photo=%.4f "
            "huber=%.4f lap=%.4f ord=%.4f. "
            "Falling back to velocity-only loss.",
            global_step,
            loss_dict["velocity"].item() if torch.is_tensor(loss_dict["velocity"]) else 0,
            loss_dict["normals"].item() if torch.is_tensor(loss_dict["normals"]) else 0,
            loss_dict["freq"].item() if torch.is_tensor(loss_dict["freq"]) else 0,
            loss_dict["grad"].item() if torch.is_tensor(loss_dict["grad"]) else 0,
            loss_dict["photo"].item() if torch.is_tensor(loss_dict["photo"]) else 0,
            loss_dict["huber"].item() if torch.is_tensor(loss_dict["huber"]) else 0,
            loss_dict["laplacian"].item() if torch.is_tensor(loss_dict["laplacian"]) else 0,
            loss_dict["ordinal"].item() if torch.is_tensor(loss_dict["ordinal"]) else 0,
        )
        total = self.vel_weight * l_vel
        if not torch.isfinite(total):
            total = torch.tensor(0.0, device=l_vel.device, requires_grad=True)
        loss_dict["total"] = total

    # Expose learned photometric parameter for logging
    if self.photo_loss is not None:
        loss_dict["lunar_lambert_weight"] = self.photo_loss.lunar_lambert_weight.detach()

    return loss_dict

depth_fm.objectives.losses¶

losses ¶

PhotoclinometricLoss ¶

surface_normals ¶

render_from_depth ¶

FlowMatchingVelocityLoss ¶

forward ¶

SurfaceNormalsLoss ¶

MultiScaleGradientLoss ¶

FocalFrequencyLoss ¶

AbsoluteDepthLoss ¶

per_pixel_loss ¶

LaplacianLoss ¶

laplacian ¶

per_pixel_loss ¶

OrdinalRankingLoss ¶

Why this helps on Mars¶

sample_pairs ¶

pair_stats ¶

CombinedLoss ¶

needs_pixel_decode ¶

forward ¶

`depth_fm.objectives.losses`¶