diff --git a/nerfactory/fields/compound_field.py b/nerfactory/fields/compound_field.py
index 49f1ac944..002d5880c 100644
--- a/nerfactory/fields/compound_field.py
+++ b/nerfactory/fields/compound_field.py
@@ -70,6 +70,7 @@ class TCNNCompoundField(Field):
         num_layers_color: number of hidden layers for color network
         hidden_dim_color: dimension of hidden layers for color network
         appearance_embedding_dim: dimension of appearance embedding
+        use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference
         spatial_distortion: spatial distortion to apply to the scene
     """
 
@@ -83,6 +84,7 @@ class TCNNCompoundField(Field):
         num_layers_color: int = 3,
         hidden_dim_color: int = 64,
         appearance_embedding_dim: int = 32,
+        use_average_appearance_embedding: bool = False,
         spatial_distortion: Optional[SpatialDistortion] = None,
     ) -> None:
         super().__init__()
@@ -94,6 +96,7 @@ class TCNNCompoundField(Field):
         self.num_images = num_images
         self.appearance_embedding_dim = appearance_embedding_dim
         self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)
+        self.use_average_appearance_embedding = use_average_appearance_embedding
 
         num_levels = 16
         max_res = 1024
@@ -171,9 +174,14 @@ class TCNNCompoundField(Field):
         if self.training:
             embedded_appearance = self.embedding_appearance(camera_indices)
         else:
-            embedded_appearance = torch.zeros(
-                (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
-            )
+            if self.use_average_appearance_embedding:
+                embedded_appearance = torch.ones(
+                    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
+                ) * self.embedding_appearance.mean(dim=0)
+            else:
+                embedded_appearance = torch.zeros(
+                    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
+                )
 
         if density_embedding is None:
             positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
diff --git a/nerfactory/fields/modules/embedding.py b/nerfactory/fields/modules/embedding.py
index f3434718f..b6c658975 100644
--- a/nerfactory/fields/modules/embedding.py
+++ b/nerfactory/fields/modules/embedding.py
@@ -41,6 +41,10 @@ class Embedding(FieldModule):
     def build_nn_modules(self) -> None:
         self.embedding = torch.nn.Embedding(self.in_dim, self.out_dim)
 
+    def mean(self, dim=0):
+        """Return the mean of the embedding weights along a dim."""
+        return self.embedding.weight.mean(dim)
+
     def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
         """Call forward
 
diff --git a/nerfactory/models/modules/ray_losses.py b/nerfactory/models/modules/ray_losses.py
index 414f1cff7..6e7255954 100644
--- a/nerfactory/models/modules/ray_losses.py
+++ b/nerfactory/models/modules/ray_losses.py
@@ -24,7 +24,6 @@ def distortion_loss(
     ray_samples: RaySamples,
     densities: TensorType["bs":..., "num_samples", 1] = None,
     weights: TensorType["bs":..., "num_samples", 1] = None,
-    scale_factor: float = 1.0,
 ) -> TensorType["bs":..., 1]:
     """Ray based distortion loss proposed in MipNeRF-360. Returns distortion Loss.
 
@@ -40,7 +39,6 @@ def distortion_loss(
         ray_samples: Ray samples to compute loss over
         densities: Predicted sample densities
         weights: Predicted weights from densities and sample locations
-        scale_factor: Scale the starts and ends by this factor
     """
     if torch.is_tensor(densities):
         assert not torch.is_tensor(weights), "Cannot use both densities and weights"
@@ -49,8 +47,8 @@ def distortion_loss(
     if torch.is_tensor(weights):
         assert not torch.is_tensor(densities), "Cannot use both densities and weights"
 
-    starts = ray_samples.frustums.starts * scale_factor
-    ends = ray_samples.frustums.ends * scale_factor
+    starts = ray_samples.spacing_starts
+    ends = ray_samples.spacing_ends
 
     midpoints = (starts + ends) / 2.0  # (..., num_samples, 1)
 
diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py
index cc1d57eb5..c27e4d891 100644
--- a/nerfactory/models/modules/ray_sampler.py
+++ b/nerfactory/models/modules/ray_sampler.py
@@ -82,11 +82,13 @@ class SpacedSampler(Sampler):
         spacing_fn_inv: Callable,
         num_samples: Optional[int] = None,
         train_stratified=True,
+        single_jitter=False,
         density_field: Optional[DensityGrid] = None,
         weight_threshold: float = 1e-2,
     ) -> None:
         super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold)
         self.train_stratified = train_stratified
+        self.single_jitter = single_jitter
         self.spacing_fn = spacing_fn
         self.spacing_fn_inv = spacing_fn_inv
 
@@ -117,7 +119,10 @@ class SpacedSampler(Sampler):
 
         # TODO More complicated than it needs to be.
         if self.train_stratified and self.training:
-            t_rand = torch.rand((num_rays, num_samples + 1), dtype=bins.dtype, device=bins.device)
+            if self.single_jitter:
+                t_rand = torch.rand((num_rays, 1), dtype=bins.dtype, device=bins.device)
+            else:
+                t_rand = torch.rand((num_rays, num_samples + 1), dtype=bins.dtype, device=bins.device)
             bin_centers = (bins[..., 1:] + bins[..., :-1]) / 2.0
             bin_upper = torch.cat([bin_centers, bins[..., -1:]], -1)
             bin_lower = torch.cat([bins[..., :1], bin_centers], -1)
@@ -152,6 +157,7 @@ class UniformSampler(SpacedSampler):
         self,
         num_samples: Optional[int] = None,
         train_stratified=True,
+        single_jitter=False,
         density_field: Optional[DensityGrid] = None,
         weight_threshold: float = 1e-2,
     ) -> None:
@@ -160,6 +166,7 @@ class UniformSampler(SpacedSampler):
             spacing_fn=lambda x: x,
             spacing_fn_inv=lambda x: x,
             train_stratified=train_stratified,
+            single_jitter=single_jitter,
             density_field=density_field,
             weight_threshold=weight_threshold,
         )
@@ -179,6 +186,7 @@ class LinearDisparitySampler(SpacedSampler):
         self,
         num_samples: Optional[int] = None,
         train_stratified=True,
+        single_jitter=False,
         density_field: Optional[DensityGrid] = None,
         weight_threshold: float = 1e-2,
     ) -> None:
@@ -187,6 +195,7 @@ class LinearDisparitySampler(SpacedSampler):
             spacing_fn=lambda x: 1 / x,
             spacing_fn_inv=lambda x: 1 / x,
             train_stratified=train_stratified,
+            single_jitter=single_jitter,
             density_field=density_field,
             weight_threshold=weight_threshold,
         )
@@ -206,6 +215,7 @@ class SqrtSampler(SpacedSampler):
         self,
         num_samples: Optional[int] = None,
         train_stratified=True,
+        single_jitter=False,
         density_field: Optional[DensityGrid] = None,
         weight_threshold: float = 1e-2,
     ) -> None:
@@ -214,6 +224,7 @@ class SqrtSampler(SpacedSampler):
             spacing_fn=torch.sqrt,
             spacing_fn_inv=lambda x: x**2,
             train_stratified=train_stratified,
+            single_jitter=single_jitter,
             density_field=density_field,
             weight_threshold=weight_threshold,
         )
@@ -233,6 +244,7 @@ class LogSampler(SpacedSampler):
         self,
         num_samples: Optional[int] = None,
         train_stratified=True,
+        single_jitter=False,
         density_field: Optional[DensityGrid] = None,
         weight_threshold: float = 1e-2,
     ) -> None:
@@ -241,6 +253,7 @@ class LogSampler(SpacedSampler):
             spacing_fn=torch.log,
             spacing_fn_inv=torch.exp,
             train_stratified=train_stratified,
+            single_jitter=single_jitter,
             density_field=density_field,
             weight_threshold=weight_threshold,
         )
@@ -262,6 +275,7 @@ class UniformLinDispPiecewiseSampler(SpacedSampler):
         self,
         num_samples: Optional[int] = None,
         train_stratified=True,
+        single_jitter=False,
         density_field: Optional[DensityGrid] = None,
         weight_threshold: float = 1e-2,
     ) -> None:
@@ -270,6 +284,7 @@ class UniformLinDispPiecewiseSampler(SpacedSampler):
             spacing_fn=lambda x: torch.where(x < 1, x / 2, 1 - 1 / (2 * x)),
             spacing_fn_inv=lambda x: torch.where(x < 0.5, 2 * x, 1 / (2 - 2 * x)),
             train_stratified=train_stratified,
+            single_jitter=single_jitter,
             density_field=density_field,
             weight_threshold=weight_threshold,
         )
@@ -531,9 +546,10 @@ class ProposalNetworkSampler(Sampler):
 
     def __init__(
         self,
-        num_proposal_samples_per_ray: int = 64,
+        num_proposal_samples_per_ray: Tuple[int] = (64,),
         num_nerf_samples_per_ray: int = 32,
         num_proposal_network_iterations: int = 2,
+        single_jitter: bool = True,
         density_field: Optional[DensityGrid] = None,
         weight_threshold: float = 1e-2,
     ) -> None:
@@ -543,7 +559,7 @@ class ProposalNetworkSampler(Sampler):
         self.num_proposal_network_iterations = num_proposal_network_iterations
 
         # samplers
-        self.initial_sampler = UniformLinDispPiecewiseSampler()
+        self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
         self.pdf_sampler = PDFSampler(include_original=False)
 
         self._anneal = 1.0
@@ -566,7 +582,7 @@ class ProposalNetworkSampler(Sampler):
         n = self.num_proposal_network_iterations
         for i_level in range(n + 1):
             is_prop = i_level < n
-            num_samples = self.num_proposal_samples_per_ray if is_prop else self.num_nerf_samples_per_ray
+            num_samples = self.num_proposal_samples_per_ray[i_level] if is_prop else self.num_nerf_samples_per_ray
             if i_level == 0:
                 # Uniform sampling because we need to start with some samples
                 ray_samples = self.initial_sampler(ray_bundle, num_samples=num_samples)
diff --git a/nerfactory/models/proposal.py b/nerfactory/models/proposal.py
index 2453552b9..7c6229a4a 100644
--- a/nerfactory/models/proposal.py
+++ b/nerfactory/models/proposal.py
@@ -34,10 +34,9 @@ from nerfactory.fields.density_field import DensityField
 from nerfactory.fields.modules.field_heads import FieldHeadNames
 from nerfactory.fields.modules.spatial_distortions import SceneContraction
 from nerfactory.models.base import Model, ModelConfig
-from nerfactory.models.modules.ray_losses import distortion_loss
 from nerfactory.models.modules.ray_sampler import ProposalNetworkSampler
 from nerfactory.models.modules.scene_colliders import NearFarCollider
-from nerfactory.optimizers.loss import MSELoss, interlevel_loss
+from nerfactory.optimizers.loss import MSELoss, distortion_loss, interlevel_loss
 from nerfactory.renderers.renderers import (
     AccumulationRenderer,
     DepthRenderer,
@@ -62,7 +61,7 @@ class ProposalModelConfig(ModelConfig):
     """How far along the ray to stop sampling."""
     randomize_background: bool = True
     """Whether to randomize the background color."""
-    num_proposal_samples_per_ray: int = 64
+    num_proposal_samples_per_ray: Tuple[int] = (64,)
     """Number of samples per ray for the proposal network."""
     num_nerf_samples_per_ray: int = 64
     """Number of samples per ray for the nerf network."""
@@ -72,16 +71,20 @@ class ProposalModelConfig(ModelConfig):
     """Use the same proposal network. Otherwise use different ones."""
     interlevel_loss_mult: float = 1.0
     """Proposal loss multiplier."""
-    distortion_loss_mult: float = 0.01
+    distortion_loss_mult: float = 0.001
     """Distortion loss multiplier."""
     use_appearance_conditioning: bool = True
     """Whether to use appearance conditioning."""
     use_proposal_weight_anneal: bool = True
     """Whether to use proposal weight annealing."""
+    use_average_appearance_embedding: bool = True
+    """Whether to use average appearance embedding or zeros for inference."""
     proposal_weights_anneal_slope: float = 10.0
     """Slope of the annealing function for the proposal weights."""
     proposal_weights_anneal_max_num_iters: int = 1000
     """Max num iterations for the annealing function."""
+    use_single_jitter: bool = False
+    """Whether use single jitter or not for first proposal network."""
 
 
 class ProposalModel(Model):
@@ -102,7 +105,10 @@ class ProposalModel(Model):
         # Fields
         if self.config.use_appearance_conditioning:
             self.field = TCNNCompoundField(
-                self.scene_bounds.aabb, spatial_distortion=scene_contraction, num_images=self.num_train_data
+                self.scene_bounds.aabb,
+                spatial_distortion=scene_contraction,
+                num_images=self.num_train_data,
+                use_average_appearance_embedding=self.config.use_average_appearance_embedding,
             )
         else:
             raise NotImplementedError("Only appearance conditioning is supported.")
@@ -127,6 +133,7 @@ class ProposalModel(Model):
             num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,
             num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,
             num_proposal_network_iterations=self.config.num_proposal_network_iterations,
+            single_jitter=self.config.use_single_jitter,
         )
 
         # renderers
@@ -197,6 +204,7 @@ class ProposalModel(Model):
         metrics_dict = {}
         image = batch["image"].to(self.device)
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
+        metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
         return metrics_dict
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
@@ -206,13 +214,7 @@ class ProposalModel(Model):
         loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
             outputs["weights_list"], outputs["ray_samples_list"]
         )
-        loss_dict["distortion_loss"] = self.config.distortion_loss_mult * torch.mean(
-            distortion_loss(
-                outputs["ray_samples_list"][-1],
-                weights=outputs["weights_list"][-1],
-                scale_factor=1.0 / (self.config.far_plane - self.config.near_plane),
-            )
-        )
+        loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"]
         return loss_dict
 
     def get_image_metrics_and_images(
diff --git a/nerfactory/optimizers/loss.py b/nerfactory/optimizers/loss.py
index 8065662d8..a25c53abb 100644
--- a/nerfactory/optimizers/loss.py
+++ b/nerfactory/optimizers/loss.py
@@ -18,6 +18,7 @@ Collection of Losses.
 
 import torch
 from torch import nn
+from torchtyping import TensorType
 
 L1Loss = nn.L1Loss
 MSELoss = nn.MSELoss
@@ -27,10 +28,12 @@ LOSSES = {"L1": L1Loss, "MSE": MSELoss}
 EPS = 1.0e-7
 
 
-def searchsorted(a, v):
+def searchsorted(a: TensorType[..., "N"], v: TensorType[..., "M"]) -> TensorType[..., "M"]:
     """
     https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L109
     https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L30
+
+
     """
     i = torch.arange(a.shape[-1], device=a.device)
     v_ge_a = v[..., None, :] >= a[..., :, None]
@@ -39,10 +42,17 @@ def searchsorted(a, v):
     return idx_lo, idx_hi
 
 
-def inner_outer(t0, t1, y1):
+def inner_outer(
+    t0: TensorType[..., "num_samples+1"], t1: TensorType[..., "num_samples+1"], y1: TensorType[..., "num_samples"]
+):
     """
     https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L117
     https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L64
+
+    Args:
+        t0: interval edges
+        t1: interval edges
+        y1: weights
     """
 
     cy1 = torch.cat([torch.zeros_like(y1[..., :1]), torch.cumsum(y1, dim=-1)], dim=-1)
@@ -61,10 +71,21 @@ def inner_outer(t0, t1, y1):
     return y0_inner, y0_outer
 
 
-def lossfun_outer(t, w, t_env, w_env):
+def lossfun_outer(
+    t: TensorType[..., "num_samples+1"],
+    w: TensorType[..., "num_samples"],
+    t_env: TensorType[..., "num_samples+1"],
+    w_env: TensorType[..., "num_samples"],
+):
     """
     https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L136
     https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L80
+
+    Args:
+        t: interval edges
+        w: weights
+        t_env: interval edges of the upper bound enveloping historgram
+        w_env: weights that should upper bound the inner (t,w) histogram
     """
     _, w_outer = inner_outer(t, t_env, w_env)
     return torch.clip(w - w_outer, min=0) ** 2 / (w + EPS)
@@ -93,3 +114,26 @@ def interlevel_loss(weights_list, ray_samples_list):
         wp = weights[..., 0]  # (num_rays, num_samples)
         loss_interlevel += torch.mean(lossfun_outer(c, w, cp, wp))
     return loss_interlevel
+
+
+# Verified
+def lossfun_distortion(t, w):
+    """
+    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L142
+    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L266
+    """
+    ut = (t[..., 1:] + t[..., :-1]) / 2
+    dut = torch.abs(ut[..., :, None] - ut[..., None, :])
+    loss_inter = torch.sum(w * torch.sum(w[..., None, :] * dut, dim=-1), dim=-1)
+
+    loss_intra = torch.sum(w**2 * (t[..., 1:] - t[..., :-1]), dim=-1) / 3
+
+    return loss_inter + loss_intra
+
+
+def distortion_loss(weights_list, ray_samples_list):
+    """From mipnerf360"""
+    c = ray_samples_to_sdist(ray_samples_list[-1])
+    w = weights_list[-1][..., 0]
+    loss = torch.mean(lossfun_distortion(c, w))
+    return loss
diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py
index 4431cdee9..435947f8d 100644
--- a/nerfactory/pipelines/base.py
+++ b/nerfactory/pipelines/base.py
@@ -131,8 +131,8 @@ class Pipeline(nn.Module):
             self.datamanager.eval_sampler.set_epoch(step)
         ray_bundle, batch = self.datamanager.next_eval(step)
         model_outputs = self.model(ray_bundle, batch)
-        loss_dict = self.model.get_loss_dict(model_outputs, batch)
         metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
+        loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
         self.train()
         return model_outputs, loss_dict, metrics_dict
 
@@ -271,8 +271,8 @@ class VanillaPipeline(Pipeline):
         if "valid_mask" in model_outputs:
             valid_mask = model_outputs["valid_mask"]
             batch = get_masked_dict(batch, valid_mask)
-        loss_dict = self.model.get_loss_dict(model_outputs, batch)
         metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
+        loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
         self.train()
         return model_outputs, loss_dict, metrics_dict