diff --git a/config/locomotion.py b/config/locomotion.py
deleted file mode 100644
index 4410bb1..0000000
--- a/config/locomotion.py
+++ /dev/null
@@ -1,70 +0,0 @@
-import socket
-
-from diffuser.utils import watch
-
-#------------------------ base ------------------------#
-
-## automatically make experiment names for planning
-## by labelling folders with these args
-
-diffusion_args_to_watch = [
-    ('prefix', ''),
-    ('horizon', 'H'),
-    ('n_diffusion_steps', 'T'),
-]
-
-base = {
-    'diffusion': {
-        ## model
-        'model': 'models.TemporalUnet',
-        'diffusion': 'models.GaussianDiffusion',
-        'horizon': 32,
-        'n_diffusion_steps': 100,
-        'action_weight': 10,
-        'loss_weights': None,
-        'loss_discount': 1,
-        'predict_epsilon': False,
-        'dim_mults': (1, 4, 8),
-        'renderer': 'utils.MuJoCoRenderer',
-
-        ## dataset
-        'loader': 'datasets.SequenceDataset',
-        'normalizer': 'LimitsNormalizer',
-        'preprocess_fns': [],
-        'clip_denoised': True,
-        'use_padding': True,
-        'max_path_length': 1000,
-
-        ## serialization
-        'logbase': 'logs',
-        'prefix': 'diffusion/',
-        'exp_name': watch(diffusion_args_to_watch),
-
-        ## training
-        'n_steps_per_epoch': 10000,
-        'loss_type': 'l2',
-        'n_train_steps': 1e6,
-        'batch_size': 32,
-        'learning_rate': 2e-4,
-        'gradient_accumulate_every': 2,
-        'ema_decay': 0.995,
-        'save_freq': 1000,
-        'sample_freq': 1000,
-        'n_saves': 5,
-        'save_parallel': False,
-        'n_reference': 8,
-        'n_samples': 2,
-        'bucket': None,
-        'device': 'cuda',
-    },
-}
-
-#------------------------ overrides ------------------------#
-
-## put environment-specific overrides here
-
-halfcheetah_medium_expert_v2 = {
-    'diffusion': {
-        'horizon': 16,
-    },
-}
diff --git a/config/maze2d.py b/config/maze2d.py
index a06ac7f..0a8d22a 100644
--- a/config/maze2d.py
+++ b/config/maze2d.py
@@ -34,11 +34,11 @@ base = {
         'model': 'models.TemporalUnet',
         'diffusion': 'models.GaussianDiffusion',
         'horizon': 256,
-        'n_diffusion_steps': 256,
+        'n_diffusion_steps': 512,
         'action_weight': 1,
         'loss_weights': None,
         'loss_discount': 1,
-        'predict_epsilon': False,
+        'predict_epsilon': True,
         'dim_mults': (1, 4, 8),
         'renderer': 'utils.Maze2dRenderer',
 
@@ -57,14 +57,14 @@ base = {
         'exp_name': watch(diffusion_args_to_watch),
 
         ## training
-        'n_steps_per_epoch': 10000,
-        'loss_type': 'l2',
-        'n_train_steps': 2e6,
-        'batch_size': 32,
-        'learning_rate': 2e-4,
-        'gradient_accumulate_every': 2,
+        'n_steps_per_epoch': 60000,
+        'loss_type': 'spline',
+        'n_train_steps': 6e4,
+        'batch_size': 1,
+        'learning_rate': 5e-6,
+        'gradient_accumulate_every': 8,
         'ema_decay': 0.995,
-        'save_freq': 1000,
+        'save_freq': 2000,
         'sample_freq': 1000,
         'n_saves': 50,
         'save_parallel': False,
@@ -89,7 +89,6 @@ base = {
         'prefix': 'plans/release',
         'exp_name': watch(plan_args_to_watch),
         'suffix': '0',
-
         'conditional': False,
 
         ## loading
@@ -122,10 +121,10 @@ maze2d_umaze_v1 = {
 maze2d_large_v1 = {
     'diffusion': {
         'horizon': 384,
-        'n_diffusion_steps': 256,
+        'n_diffusion_steps': 16,
     },
     'plan': {
         'horizon': 384,
-        'n_diffusion_steps': 256,
+        'n_diffusion_steps': 16,
     },
 }
diff --git a/diffuser/datasets/buffer.py b/diffuser/datasets/buffer.py
index 1ad2106..5991f01 100644
--- a/diffuser/datasets/buffer.py
+++ b/diffuser/datasets/buffer.py
@@ -9,7 +9,7 @@ class ReplayBuffer:
 
     def __init__(self, max_n_episodes, max_path_length, termination_penalty):
         self._dict = {
-            'path_lengths': np.zeros(max_n_episodes, dtype=np.int),
+            'path_lengths': np.zeros(max_n_episodes, dtype=np.int_),
         }
         self._count = 0
         self.max_n_episodes = max_n_episodes
diff --git a/diffuser/datasets/sequence.py b/diffuser/datasets/sequence.py
index 356c540..73c1b04 100644
--- a/diffuser/datasets/sequence.py
+++ b/diffuser/datasets/sequence.py
@@ -83,6 +83,7 @@ class SequenceDataset(torch.utils.data.Dataset):
         actions = self.fields.normed_actions[path_ind, start:end]
 
         conditions = self.get_conditions(observations)
+
         trajectories = np.concatenate([actions, observations], axis=-1)
         batch = Batch(trajectories, conditions)
         return batch
diff --git a/diffuser/models/diffusion.py b/diffuser/models/diffusion.py
index fae4cfd..461680a 100644
--- a/diffuser/models/diffusion.py
+++ b/diffuser/models/diffusion.py
@@ -2,6 +2,7 @@ import numpy as np
 import torch
 from torch import nn
 import pdb
+import matplotlib.pyplot as plt
 
 import diffuser.utils as utils
 from .helpers import (
@@ -9,6 +10,7 @@ from .helpers import (
     extract,
     apply_conditioning,
     Losses,
+    catmull_rom_spline_with_rotation,
 )
 
 class GaussianDiffusion(nn.Module):
@@ -26,6 +28,7 @@ class GaussianDiffusion(nn.Module):
         betas = cosine_beta_schedule(n_timesteps)
         alphas = 1. - betas
         alphas_cumprod = torch.cumprod(alphas, axis=0)
+        print(f"Alphas Cumprod: {alphas_cumprod}")
         alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])
 
         self.n_timesteps = int(n_timesteps)
@@ -73,7 +76,7 @@ class GaussianDiffusion(nn.Module):
         '''
         self.action_weight = action_weight
 
-        dim_weights = torch.ones(self.transition_dim, dtype=torch.float32)
+        dim_weights = torch.ones(self.transition_dim, dtype=torch.float64)
 
         ## set loss coefficients for dimensions of observation
         if weights_dict is None: weights_dict = {}
@@ -97,18 +100,16 @@ class GaussianDiffusion(nn.Module):
             otherwise, model predicts x0 directly
         '''
         if self.predict_epsilon:
-            return (
-                extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
-                extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
-            )
+            return noise            
         else:
             return noise
 
     def q_posterior(self, x_start, x_t, t):
         posterior_mean = (
             extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
-            extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+            extract(self.posterior_mean_coef2, t, x_t.shape) * x_t[:, :, self.action_dim:]
         )
+
         posterior_variance = extract(self.posterior_variance, t, x_t.shape)
         posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
@@ -129,7 +130,7 @@ class GaussianDiffusion(nn.Module):
     def p_sample(self, x, cond, t):
         b, *_, device = *x.shape, x.device
         model_mean, _, model_log_variance = self.p_mean_variance(x=x, cond=cond, t=t)
-        noise = torch.randn_like(x)
+        noise = torch.randn_like(x[:, :, self.action_dim:])
         # no noise when t == 0
         nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
@@ -139,22 +140,59 @@ class GaussianDiffusion(nn.Module):
         device = self.betas.device
 
         batch_size = shape[0]
-        x = torch.randn(shape, device=device)
-        x = apply_conditioning(x, cond, self.action_dim)
+        # x = torch.randn(shape, device=device, dtype=torch.float64)
+        # Extract known indices and values
+        known_indices = np.array(list(cond.keys()), dtype=int)
+
+        # candidate_no x batch_size x dim
+        known_values = np.stack([c.cpu().numpy() for c in cond.values()], axis=0)
+        known_values = np.moveaxis(known_values, 0, 1)
+
+        # Sort the timepoints
+        sorted_indices = np.argsort(known_indices)
+        known_indices = known_indices[sorted_indices]
+        known_values = known_values[:, sorted_indices]
+
+        # Build the structured spline guess
+        catmull_spline_trajectory = np.array([
+            catmull_rom_spline_with_rotation(known_values[b, :, :-1], known_indices, shape[1]) 
+            for b in range(batch_size)
+        ])
+        catmull_spline_trajectory = torch.tensor(
+            catmull_spline_trajectory, 
+            dtype=torch.float64, 
+            device=device
+        )            
+
+
+        if self.predict_epsilon:
+            x = torch.randn((shape[0], shape[1], self.observation_dim), device=device, dtype=torch.float64)
+            cond_residual = {k: torch.zeros_like(v)[:, :-1] for k, v in cond.items()}
+            is_cond = torch.zeros((shape[0], shape[1], 1), device=device, dtype=torch.float64)
+            is_cond[:, known_indices, :] = 1.0             
 
         if return_diffusion: diffusion = [x]
 
-        progress = utils.Progress(self.n_timesteps) if verbose else utils.Silent()
+        # progress = utils.Progress(self.n_timesteps) if verbose else utils.Silent()
         for i in reversed(range(0, self.n_timesteps)):
+            if self.predict_epsilon:
+                x = torch.cat([catmull_spline_trajectory, is_cond, x], dim=-1)
+            
             timesteps = torch.full((batch_size,), i, device=device, dtype=torch.long)
-            x = self.p_sample(x, cond, timesteps)
-            x = apply_conditioning(x, cond, self.action_dim)
+            x = self.p_sample(x, cond_residual, timesteps)
+
+            x = apply_conditioning(x, cond_residual, 0)
 
-            progress.update({'t': i})
+            if return_diffusion: diffusion.append(x)        
 
-            if return_diffusion: diffusion.append(x)
+        x = catmull_spline_trajectory + x
 
-        progress.close()
+        
+        
+        # Normalize the quaternions
+        # x[:, :, 3:7] = x[:, :, 3:7] / torch.norm(x[:, :, 3:7], dim=-1, keepdim=True)
+
+        # progress.close()
 
         if return_diffusion:
             return x, torch.stack(diffusion, dim=1)
@@ -167,7 +205,7 @@ class GaussianDiffusion(nn.Module):
             conditions : [ (time, state), ... ]
         '''
         device = self.betas.device
-        batch_size = len(cond[0])
+        batch_size = len(next(iter(cond.values())))
         horizon = horizon or self.horizon
         shape = (batch_size, horizon, self.transition_dim)
 
@@ -175,38 +213,106 @@ class GaussianDiffusion(nn.Module):
 
     #------------------------------------------ training ------------------------------------------#
 
-    def q_sample(self, x_start, t, noise=None):
+    def q_sample(self, x_start, t, spline=None, noise=None):
+        x_start_noise = x_start[:, : , :-1]
+        x_start_is_cond = x_start[:, :, [-1]]
+
+        if spline is None:
+            spline = torch.randn_like(x_start_noise)
         if noise is None:
-            noise = torch.randn_like(x_start)
+            noise = torch.randn_like(x_start_noise)
 
-        sample = (
-            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
-            extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
-        )
+        alpha = extract(self.sqrt_alphas_cumprod, t, x_start.shape)
+        oneminusalpha = extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+
+        # Weighted combination of x_0 and the spline
+        out = alpha * x_start_noise + oneminusalpha * noise
+
+        # Concatenate the binary feature and the spline as the conditioning
+        out = torch.cat([spline, x_start_is_cond, out], dim=-1)
 
-        return sample
+        return out
 
     def p_losses(self, x_start, cond, t):
-        noise = torch.randn_like(x_start)
+        batch_size, horizon, _ = x_start.shape
+        # Extract known indices and values
+        known_indices = np.array(list(cond.keys()), dtype=int)
+
+        # candidate_no x batch_size x dim
+        known_values = np.stack([c.cpu().numpy() for c in cond.values()], axis=0)
+        known_values = np.moveaxis(known_values, 0, 1)
+
+        # Sort the timepoints
+        sorted_indices = np.argsort(known_indices)
+        known_indices = known_indices[sorted_indices]
+        known_values = known_values[:, sorted_indices]
+
+        # Build your structured guess
+        catmull_spline_trajectory = np.array([
+            catmull_rom_spline_with_rotation(known_values[b, :, :-1], known_indices, horizon) 
+            for b in range(batch_size)
+        ])
+        catmull_spline_trajectory = torch.tensor(
+            catmull_spline_trajectory, 
+            dtype=torch.float64, 
+            device=x_start.device
+        )
 
-        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
-        x_noisy = apply_conditioning(x_noisy, cond, self.action_dim)
+        # Plot the quaternions
+        # plt.plot(x_start[0, :, 3].cpu().numpy())
+        # plt.plot(catmull_spline_trajectory[0, :, 3].cpu().numpy())
+        # plt.legend(["x_start", "catmull_spline"])
+        # plt.show()
+        # raise Exception    
 
-        x_recon = self.model(x_noisy, cond, t)
-        x_recon = apply_conditioning(x_recon, cond, self.action_dim)
 
-        assert noise.shape == x_recon.shape
+        if not self.predict_epsilon:
+            # Forward diffuse with the structured trajectory
+            x_noisy = self.q_sample(
+                x_start,
+                t,
+                spline=catmull_spline_trajectory,
+            )
+            x_noisy = apply_conditioning(x_noisy, cond, self.action_dim)
 
-        if self.predict_epsilon:
-            loss, info = self.loss_fn(x_recon, noise)
+            # Reverse pass guess
+            x_recon = self.model(x_noisy, cond, t)
+            x_recon = apply_conditioning(x_recon, cond, self.action_dim)
+
+            # Then x_recon is the predicted x_0, compare to the true x_0
+            loss, info = self.loss_fn(x_recon, x_start, cond)
         else:
-            loss, info = self.loss_fn(x_recon, x_start)
+            residual = x_start.clone()
+
+            residual[:, :, :-1] -= catmull_spline_trajectory
+            
+
+            cond_residual = {k: torch.zeros_like(v)[:, :-1] for k, v in cond.items()}
+
+            x_noisy = self.q_sample(
+                residual,
+                t,
+                spline=catmull_spline_trajectory,
+            )
+            x_noisy = apply_conditioning(x_noisy, cond_residual, self.action_dim)
+
+            # Reverse pass guess
+            x_recon = self.model(x_noisy, cond, t)
+            x_recon = apply_conditioning(x_recon, cond_residual, 0)
+
+            x_recon = x_recon + catmull_spline_trajectory
+
+            loss, info = self.loss_fn(x_recon, x_start[:, :, :-1], cond)
 
         return loss, info
 
     def loss(self, x, cond):
         batch_size = len(x)
         t = torch.randint(0, self.n_timesteps, (batch_size,), device=x.device).long()
+        # t = torch.randint(1, 2, (batch_size,), device=x.device).long()
+        # x = x.double()
+        # cond = {k: v.double() for k, v in cond.items()}
+        # print(f"Time: {t.item()}")
         return self.p_losses(x, cond, t)
 
     def forward(self, cond, *args, **kwargs):
diff --git a/diffuser/models/helpers.py b/diffuser/models/helpers.py
index d39f35d..9f43ef8 100644
--- a/diffuser/models/helpers.py
+++ b/diffuser/models/helpers.py
@@ -1,11 +1,11 @@
 import math
+import json
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import einops
 from einops.layers.torch import Rearrange
-import pdb
+from pytorch3d.transforms import quaternion_to_matrix, quaternion_to_axis_angle
 
 import diffuser.utils as utils
 
@@ -30,7 +30,7 @@ class SinusoidalPosEmb(nn.Module):
 class Downsample1d(nn.Module):
     def __init__(self, dim):
         super().__init__()
-        self.conv = nn.Conv1d(dim, dim, 3, 2, 1)
+        self.conv = nn.Conv1d(dim, dim, 3, 2, 1).to(torch.float64)
 
     def forward(self, x):
         return self.conv(x)
@@ -38,7 +38,7 @@ class Downsample1d(nn.Module):
 class Upsample1d(nn.Module):
     def __init__(self, dim):
         super().__init__()
-        self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1)
+        self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1).to(torch.float64)
 
     def forward(self, x):
         return self.conv(x)
@@ -52,9 +52,9 @@ class Conv1dBlock(nn.Module):
         super().__init__()
 
         self.block = nn.Sequential(
-            nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2),
+            nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2).to(torch.float64),
             Rearrange('batch channels horizon -> batch channels 1 horizon'),
-            nn.GroupNorm(n_groups, out_channels),
+            nn.GroupNorm(n_groups, out_channels).to(torch.float64),
             Rearrange('batch channels 1 horizon -> batch channels horizon'),
             nn.Mish(),
         )
@@ -72,7 +72,7 @@ def extract(a, t, x_shape):
     out = a.gather(-1, t)
     return out.reshape(b, *((1,) * (len(x_shape) - 1)))
 
-def cosine_beta_schedule(timesteps, s=0.008, dtype=torch.float32):
+def cosine_beta_schedule(timesteps, s=0.008, dtype=torch.float64):
     """
     cosine schedule
     as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
@@ -157,9 +157,979 @@ class ValueL2(ValueLoss):
     def _loss(self, pred, targ):
         return F.mse_loss(pred, targ, reduction='none')
 
+class GeodesicL2Loss(nn.Module):
+    def __init__(self, *args):
+        super().__init__()
+        pass
+    
+    def _loss(self, pred, targ):
+        # Compute L2 loss for the first three dimensions
+        l2_loss = F.mse_loss(pred[..., :3], targ[..., :3], reduction='mean')
+        
+        # Normalize to unit quaternions for the last four dimensions
+        pred_quat = pred[..., 3:] / pred[..., 3:].norm(dim=-1, keepdim=True)
+        targ_quat = targ[..., 3:] / targ[..., 3:].norm(dim=-1, keepdim=True)
+
+        assert not torch.isnan(pred_quat).any(), "Pred Quat has NaNs"
+        assert not torch.isnan(targ_quat).any(), "Targ Quat has NaNs"
+        
+        # Compute dot product for the quaternions
+        dot_product = torch.sum(pred_quat * targ_quat, dim=-1)
+        dot_product = torch.clamp(torch.abs(dot_product), -1.0, 1.0)
+        
+        # Compute geodesic loss for the quaternions
+        geodesic_loss = 2 * torch.acos(dot_product).mean()
+
+        assert not torch.isnan(geodesic_loss).any(), "Geodesic Loss has NaNs"
+        assert not torch.isnan(l2_loss).any(), "L2 Loss has NaNs"
+
+        return l2_loss + geodesic_loss, l2_loss, geodesic_loss
+
+    def forward(self, pred, targ):
+        loss, l2, geodesic = self._loss(pred, targ)
+        
+        info = {
+            'l2': l2.item(),
+            'geodesic': geodesic.item(),
+        }
+        
+        return loss, info
+
+class RotationTranslationLoss(nn.Module):
+    def __init__(self, *args):
+        super().__init__()
+        pass
+    
+    def _loss(self, pred, targ, cond=None):
+        
+        # Make sure the dtype is float64
+        pred = pred.to(torch.float64)
+        targ = targ.to(torch.float64)
+        
+        eps = 1e-8
+
+        pred_trans = pred[..., :3]
+        pred_quat = pred[..., 3:7]
+        targ_trans = targ[..., :3]
+        targ_quat = targ[..., 3:7]
+
+        l2_loss = F.mse_loss(pred_trans, targ_trans, reduction='mean')
+
+        # Calculate the geodesic loss
+        pred_n = pred_quat.norm(dim=-1, keepdim=True).clamp(min=eps)
+        targ_n = targ_quat.norm(dim=-1, keepdim=True).clamp(min=eps)
+        
+        pred_quat_norm = pred_quat / pred_n
+        targ_quat_norm = targ_quat / targ_n
+
+
+        dot_product = torch.sum(pred_quat_norm * targ_quat_norm, dim=-1).clamp(min=-1.0 + eps, max=1.0 - eps)
+        quaternion_dist = 1 - (dot_product ** 2).mean()
+
+        # Calculate the rotation error
+        pred_rot = quaternion_to_matrix(pred_quat_norm).reshape(-1, 3, 3)
+        targ_rot = quaternion_to_matrix(targ_quat_norm).reshape(-1, 3, 3)
+
+        r2r1 = pred_rot @ targ_rot.permute(0, 2, 1)
+        trace = torch.diagonal(r2r1, dim1=-2, dim2=-1).sum(-1)
+        trace = torch.clamp((trace - 1) / 2, -1.0 + eps, 1.0 - eps)
+        geodesic_loss = torch.acos(trace).mean()
+
+        # Add a smoothness and acceleration term to the positions and quaternions
+        alpha = 1.0
+        smoothness_loss = F.mse_loss(pred[:, 1:, :7].reshape(-1, 7), pred[:, :-1, :7].reshape(-1, 7), reduction='mean')
+        acceleration_loss = F.mse_loss(pred[:, 2:, :7].reshape(-1, 7), 2 * pred[:, 1:-1, :7].reshape(-1, 7) - pred[:, :-2, :7].reshape(-1, 7), reduction='mean')
+
+        l2_multiplier = 10.0
+
+        loss = l2_multiplier * l2_loss + quaternion_dist + geodesic_loss + alpha * (smoothness_loss + acceleration_loss)
+
+        dtw = DynamicTimeWarpingLoss()
+        dtw_loss, _ = dtw.forward(pred_trans.reshape(-1, 3), targ_trans.reshape(-1, 3))
+
+        hausdorff = HausdorffDistanceLoss()
+        hausdorff_loss, _ = hausdorff.forward(pred_trans.reshape(-1, 3), targ_trans.reshape(-1, 3))
+
+        frec = FrechetDistanceLoss()
+        frechet_loss, _ = frec.forward(pred_trans.reshape(-1, 3), targ_trans.reshape(-1, 3))
+
+        chamfer = ChamferDistanceLoss()
+        chamfer_loss, _ = chamfer.forward(pred_trans.reshape(-1, 3), targ_trans.reshape(-1, 3))
+
+        return loss, l2_loss, geodesic_loss, quaternion_dist, dtw_loss, hausdorff_loss, frechet_loss, chamfer_loss
+
+
+    def forward(self, pred, targ, cond=None):
+        loss, err_t, err_geo, err_r, err_dtw, err_hausdorff, err_frechet, err_chamfer = self._loss(pred, targ, cond)
+        
+        info = {
+            'rot. error': err_r.item(),
+            'geodesic error': err_geo.item(),
+            'trans. error': err_t.item(),
+            'dtw': err_dtw.item(),
+            'hausdorff': err_hausdorff.item(),
+            'frechet': err_frechet.item(),
+            'chamfer': err_chamfer.item(),
+        }
+        
+        return loss, info
+
+class SplineLoss(nn.Module):
+    def __init__(self, *args):
+        super().__init__()
+        self.scales = json.load(open('scene_scale.json'))
+
+    def compute_spline_coeffs(self, trans):
+        p0 = trans[:, :-3, :]
+        p1 = trans[:, 1:-2, :]
+        p2 = trans[:, 2:-1, :]
+        p3 = trans[:, 3:, :]
+
+        # Tangent approximations
+        m1 = 0.5 * (-p0 + p2)
+        m2 = 0.5 * (-p1 + p3)
+
+        # Cubic spline coefficients for each dimension
+        a = (2 * p1 - 2 * p2 + m1 + m2)
+        b = (-3 * p1 + 3 * p2 - 2 * m1 - m2)
+        c = (m1)
+        d = (p1)
+
+        return torch.stack([a, b, c, d], dim=-1)
+    
+    def q_normalize(self, q):
+        return q / q.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-12)
+    
+    def q_conjugate(self, q):
+        w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
+        return torch.stack([w, -x, -y, -z], dim=-1)
+
+    def q_multiply(self, q1, q2):
+        """
+        q1*q2.
+        """
+        w1, x1, y1, z1 = q1.unbind(-1)
+        w2, x2, y2, z2 = q2.unbind(-1)
+        w = w1*w2 - x1*x2 - y1*y2 - z1*z2
+        x = w1*x2 + x1*w2 + y1*z2 - z1*y2
+        y = w1*y2 - x1*z2 + y1*w2 + z1*x2
+        z = w1*z2 + x1*y2 - y1*x2 + z1*w2
+        return torch.stack([w, x, y, z], dim=-1)
+
+    def q_inverse(self, q):
+        return self.q_conjugate(self.q_normalize(q))
+
+    def q_log(self, q):
+        """
+        Quaternion logarithm for a unit quaternion
+        Only returns the imaginary part
+        """
+        q = self.q_normalize(q)
+        w = q[..., 0]
+        xyz = q[..., 1:]  # shape [..., 3]
+        mag_v = xyz.norm(p=2, dim=-1)
+        eps = 1e-12
+        angle = torch.acos(w.clamp(-1.0 + eps, 1.0 - eps))
+
+        # We do a safe-guard against zero for sin(angle)
+        small_mask = (mag_v < 1e-12) | (angle < 1e-12)
+        # Where small_mask is True => near identity => log(q) ~ 0
+        log_val = torch.zeros_like(xyz)
+
+        # Normal case
+        scale = angle / mag_v.clamp(min=1e-12)
+        normal_case = scale.unsqueeze(-1) * xyz
+
+        log_val = torch.where(
+            small_mask.unsqueeze(-1),
+            torch.zeros_like(xyz),
+            normal_case
+        )
+        return log_val
+
+    def q_exp(self, v):
+        """
+        Quaternion exponential
+        """
+        norm_v = v.norm(p=2, dim=-1)
+        small_mask = norm_v < 1e-12
+
+        w = torch.cos(norm_v)
+        sin_v = torch.sin(norm_v)
+        scale = torch.where(
+            small_mask,
+            torch.zeros_like(norm_v),  # if zero, sin(0)/0 => 0
+            sin_v / norm_v.clamp(min=1e-12)
+        )
+        xyz = scale.unsqueeze(-1) * v
+
+        # For small angles, we approximate cos(norm_v) ~ 1, sin(norm_v)/norm_v ~ 1
+        w = torch.where(
+            small_mask,
+            torch.ones_like(w),
+            w
+        )
+        return torch.cat([w.unsqueeze(-1), xyz], dim=-1)
+
+    def q_slerp(self, q1, q2, t):
+        """
+        Spherical linear interpolation from q1 to q2 at t in [0,1].
+        Both q1, q2 assumed normalized.
+        q1, q2, t can be 1D or broadcastable shapes, but typically 1D.
+        """
+        q1 = self.q_normalize(q1)
+        q2 = self.q_normalize(q2)
+        dot = (q1 * q2).sum(dim=-1, keepdim=True)  # the dot product
+
+        eps = 1e-12
+        dot = dot.clamp(-1.0 + eps, 1.0 - eps)
+
+        flip_mask = dot < 0.0
+        if flip_mask.any():
+            q2 = torch.where(flip_mask, -q2, q2)
+            dot = torch.where(flip_mask, -dot, dot)
+
+        # If they're very close, do a simple linear interpolation
+        close_mask = dot.squeeze(-1) > 0.9995
+        # Using an epsilon to avoid potential issues close to 1.0
+
+        # Branch 1: Very close
+        # linear LERP
+        lerp_val = (1.0 - t) * q1 + t * q2
+        lerp_val = self.q_normalize(lerp_val)
+
+        # Branch 2: Standard SLERP
+        theta_0 = torch.acos(dot)
+        sin_theta_0 = torch.sin(theta_0)
+        theta = theta_0 * t
+        s1 = torch.sin(theta_0 - theta) / sin_theta_0.clamp(min=1e-12)
+        s2 = torch.sin(theta) / sin_theta_0.clamp(min=1e-12)
+        slerp_val = s1 * q1 + s2 * q2
+        slerp_val = self.q_normalize(slerp_val)
+
+        # Combine
+        return torch.where(
+            close_mask.unsqueeze(-1),
+            lerp_val,
+            slerp_val
+        )
+
+    def compute_uniform_tangent(self, q_im1, q_i, q_ip1):
+        """
+        Computes a 'Catmull–Rom-like' tangent T_i for quaternion q_i,
+        given neighbors q_im1, q_i, q_ip1.
+
+        T_i = q_i * exp( -0.25 * [ log(q_i^-1 q_ip1) + log(q_i^-1 q_im1) ] )
+        """
+        q_im1 = self.q_normalize(q_im1)
+        q_i   = self.q_normalize(q_i)
+        q_ip1 = self.q_normalize(q_ip1)
+
+        inv_qi = self.q_inverse(q_i)
+        r1 = self.q_multiply(inv_qi, q_ip1)
+        r2 = self.q_multiply(inv_qi, q_im1)
+
+        lr1 = self.q_log(r1)
+        lr2 = self.q_log(r2)
+
+        m = -0.25 * (lr1 + lr2)
+        exp_m = self.q_exp(m)
+        return self.q_multiply(q_i, exp_m)
+
+    def compute_all_uniform_tangents(self, quats):
+        """
+        Vectorized version that computes tangents T_i for all keyframe quaternions at once.
+        quats shape: [N,4], N >= 2
+        Returns shape [N,4].
+        """
+        q_im1 = torch.cat([quats[[0]], quats[:-1]], dim=0)   # q_im1[0]   = q0
+        q_ip1 = torch.cat([quats[1:], quats[[-1]]], dim=0)   # q_ip1[N-1]= q_{N-1}
+
+        return self.compute_uniform_tangent(q_im1, quats, q_ip1)
+
+    def squad(self, q0, a, b, q1, t):
+        """
+        Shoemake's "squad" interpolation for quaternion splines:
+            squad(q0, a, b, q1; t) = slerp( slerp(q0, q1; t),
+                                            slerp(a,   b;   t),
+                                            2t(1-t) )
+        where a, b are tangential control quaternions for q0, q1.
+        """
+        s1 = self.q_slerp(q0, q1, t)
+        s2 = self.q_slerp(a,   b,   t)
+        alpha = 2.0*t*(1.0 - t)
+        return self.q_slerp(s1, s2, alpha)
+    
+    def uniform_cr_spline(self, quats, num_samples_per_segment=10):
+        """
+        Given a list of keyframe quaternions quats (each a torch 1D tensor [4]),
+        compute a "Uniform Catmull–Rom–like" quaternion spline through them.
+
+        Returns:
+          A list (Python list) of interpolated quaternions (torch tensors),
+          including all segment endpoints.
+
+        Each interior qi gets a tangent T_i using neighbors q_{i-1}, q_i, q_{i+1}.
+        For boundary tangents, we replicate the end quaternions.
+        """
+        n = quats.shape[0]
+        if n < 2:
+            return quats.unsqueeze(0)  # not enough quats to interpolate
+
+        # Precompute tangents
+        tangents = self.compute_all_uniform_tangents(quats)
+
+        # Interpolate each segment [qi, q_{i+1}]
+        q0 = quats[:-1].unsqueeze(1) 
+        q1 = quats[1:].unsqueeze(1)  
+        a = tangents[:-1].unsqueeze(1) 
+        b = tangents[1:].unsqueeze(1)  
+
+        t_vals = torch.linspace(0.0, 1.0, num_samples_per_segment, device=quats.device, dtype=quats.dtype)
+        t_vals = t_vals.view(1, -1, 1)
+
+        out = self.squad(q0, a, b, q1, t_vals)
+        return out
+
+
+    def forward(self, pred, targ, cond=None, scene_id=None, norm_params=None):
+        loss, err_t, err_smooth, err_geo, err_r, err_dtw, err_hausdorff, err_frechet, err_chamfer = self._loss(pred, targ, cond, scene_id, norm_params)
+        
+        info = {
+            'trans. error': err_t.item(),
+            'smoothness error': err_smooth.item(),
+            # 'dtw': err_dtw.item(),
+            # 'hausdorff': err_hausdorff.item(),
+            # 'frechet': err_frechet.item(),
+            # 'chamfer': err_chamfer.item(),
+            'quat. dist.': err_r.item(),
+            'geodesic dist.': err_geo.item(),
+        }
+        
+        return loss, info
+
+    def _loss(self, pred, targ, cond=None, scene_id=None, norm_params=None):
+        def poly_eval(coeffs, x):
+            """
+            Evaluates a polynomial (with highest-degree term first) at points x.
+            coeffs: 2D tensor of shape [num_polynomials, degree + 1], highest-degree term first.
+            x: 1D tensor of points at which to evaluate the polynomial.
+            Returns:
+            2D tensor of shape [num_polynomials, len(x)], containing p(x).
+            """
+            x_powers = torch.stack([x**i for i in range(coeffs.shape[-1] - 1, -1, -1)], dim=-1)
+            x_powers = x_powers.to(torch.float64).to(coeffs.device)
+            y = torch.matmul(coeffs, x_powers.T)
+            return y
+        
+        # Make sure the dtype is float64
+        pred = pred.to(torch.float64)
+        targ = targ.to(torch.float64)
+
+        # Rescale the translations
+        if scene_id is not None and norm_params is not None:
+            scene_id = scene_id.item()
+            scene_scale = self.scales[str(scene_id)]
+            scene_scale = norm_params['scale'][0] * scene_scale
+            pred[..., :3] = pred[..., :3] * scene_scale
+            targ[..., :3] = targ[..., :3] * scene_scale
+            # print(pred[..., :3].max(), targ[..., :3].max())
+
+        # We only consider interpolated points for loss calculation
+        candidate_idxs = sorted(cond.keys())
+        pred = pred[:, candidate_idxs[0] : candidate_idxs[-1] + 1, :]
+        targ = targ[:, candidate_idxs[0] : candidate_idxs[-1] + 1, :]
+        
+        pred_trans = pred[..., :3]
+        pred_quat = pred[..., 3:7]
+        targ_trans = targ[..., :3]
+        targ_quat = targ[..., 3:7]
+
+        pred_coeffs = self.compute_spline_coeffs(pred_trans)
+        targ_coeffs = self.compute_spline_coeffs(targ_trans)
+        
+        n_points = 2000
+
+        # Distribute sample points among intervals
+        dists = torch.norm(targ_trans[:, 1:, :] - targ_trans[:, :-1, :], dim=-1).reshape(-1)
+        dists_c = torch.zeros(len(candidate_idxs) - 1, device=pred.device)
+        for i in range(len(candidate_idxs) - 1):
+            dists_c[i] = dists[candidate_idxs[i]:candidate_idxs[i+1]].sum()
+        
+        weights_c = dists_c / dists_c.sum()
+        scaled_c = weights_c * n_points
+        points_c = torch.floor(scaled_c).int()
+
+        while points_c.sum() < n_points:
+            idx = torch.argmax(scaled_c - points_c)
+            points_c[idx] += 1 
+        
+        # Calculate the spline loss
+        sample_points = 50
+        x = torch.linspace(0, 1, sample_points, device=pred.device)
+        pred_spline = poly_eval(pred_coeffs, x).permute(0, 1, 3, 2).reshape(-1, sample_points, 3)
+        targ_spline = poly_eval(targ_coeffs, x).permute(0, 1, 3, 2).reshape(-1, sample_points, 3)
+
+        indexes = []
+        start_idx = candidate_idxs[0]
+        for c, (idx_i0, idx_i1) in enumerate(zip(candidate_idxs[:-1], candidate_idxs[1:])):
+            p = points_c[c]
+            total_dist = dists_c[c]
+            dist_arr = dists[idx_i0 - start_idx : idx_i1 - start_idx]
+            
+            step_distances = (dist_arr / sample_points).repeat_interleave(sample_points)
+            cumul_distances = step_distances.cumsum(dim=0)
+
+            dist_per_pick = total_dist / p
+            pick_targets = torch.arange(1, p + 1, device=dists.device) * dist_per_pick
+
+            pick_idxs = torch.searchsorted(cumul_distances, pick_targets, right=True)
+            pick_idxs = torch.clamp(pick_idxs, max=len(cumul_distances) - 1)
+
+
+            indexes_1d = torch.zeros_like(step_distances)
+            indexes_1d[pick_idxs] = 1
+
+            indexes_2d = indexes_1d.view(len(dist_arr), sample_points)
+
+            indexes.append(indexes_2d)
+
+        indexes = torch.cat(indexes)[1: -1] # The first and last candidates don't have spline representations
+     
+        indexes_trans = torch.stack([indexes for _ in range(3)], dim=-1)
+        indexes_quat = torch.stack([indexes for _ in range(4)], dim=-1)
+
+        indexes_trans = indexes_trans.to(torch.bool)
+        indexes_quat = indexes_quat.to(torch.bool)
+
+        pred_trans_selected_values = pred_spline[indexes_trans]
+        targ_trans_selected_values = targ_spline[indexes_trans]
+
+        pred_trans_selected_values = pred_trans_selected_values.reshape(-1, 3)
+        targ_trans_selected_values = targ_trans_selected_values.reshape(-1, 3)
+
+        # Calculate the loss for quaternions
+        pred_quat = pred_quat / pred_quat.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+        targ_quat = targ_quat / targ_quat.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+
+        targ_quat_spline = self.uniform_cr_spline(targ_quat.reshape(-1, 4), num_samples_per_segment=sample_points)
+        pred_quat_spline = self.uniform_cr_spline(pred_quat.reshape(-1, 4), num_samples_per_segment=sample_points)
+
+        
+        targ_quat_spline = targ_quat_spline[1:-1]
+        pred_quat_spline = pred_quat_spline[1:-1]
+
+
+        pred_quat_selected_values = pred_quat_spline[indexes_quat]
+        targ_quat_selected_values = targ_quat_spline[indexes_quat]
+
+        pred_quat_selected_values = pred_quat_selected_values.reshape(-1, 4)
+        targ_quat_selected_values = targ_quat_selected_values.reshape(-1, 4)
+
+        # Calculate the geodesic loss
+        pred_rot = quaternion_to_matrix(pred_quat_selected_values).reshape(-1, 3, 3)
+        targ_rot = quaternion_to_matrix(targ_quat_selected_values).reshape(-1, 3, 3)
+
+        eps = 1e-12
+        r2r1 = pred_rot @ targ_rot.permute(0, 2, 1)
+        trace = torch.diagonal(r2r1, dim1=-2, dim2=-1).sum(-1)
+        trace = torch.clamp((trace - 1) / 2, -1.0 + eps, 1.0 - eps)
+        geodesic_loss = torch.acos(trace).mean()
+
+        # Calculate the rotation error
+        dot_product = torch.sum(pred_quat_selected_values * targ_quat_selected_values, dim=-1).clamp(min=-1.0 + eps, max=1.0 - eps)
+        quaternion_dist = 1 - (dot_product ** 2).mean()
+
+        # Calculate the L2 loss
+        l2_loss = F.mse_loss(pred_trans_selected_values, targ_trans_selected_values, reduction='mean')
+
+        # Calculate the smoothness loss for translation and quaternion
+        smoothness_multiplier = 10 ** 2 # Empirically determined multiplier for smoothness loss
+        weight_acceleration = 0.1 
+        weight_jerk = 0.05 
+
+        pos_acc = pred_trans_selected_values[2:, :] - 2 * pred_trans_selected_values[1:-1, :] + pred_trans_selected_values[:-2, :]
+        pos_jerk = pred_trans_selected_values[3:, :] - 3 * pred_trans_selected_values[2:-1, :] + 3 * pred_trans_selected_values[1:-2, :] - pred_trans_selected_values[:-3, :]
+
+        pos_acceleration_loss = torch.mean(pos_acc ** 2)
+        pos_jerk_loss = torch.mean(pos_jerk ** 2)
+
+        q0 = pred_quat_selected_values[:-1, :]
+        q1 = pred_quat_selected_values[1:, :]
+        sign = torch.where((q0 * q1).sum(dim=-1) < 0, -1.0, 1.0)
+        q1 = sign.unsqueeze(-1) * q1
+
+        dq = self.q_multiply(q1, self.q_inverse(q0))
+        theta = 2 * torch.acos(torch.clamp(dq[..., 0], -1.0 + 1e-8, 1.0 - 1e-8))
+
+        rot_acc  = theta[2:] - 2*theta[1:-1] + theta[:-2]
+        rot_jerk = theta[3:] - 3*theta[2:-1] + 3*theta[1:-2] - theta[:-3] 
+
+        rot_acceleration_loss = torch.mean(rot_acc ** 2)
+        rot_jerk_loss = torch.mean(rot_jerk ** 2)
+
+        alpha_rot = 0.1        # <-- tune this (e.g. 0.1 … 10)
+
+
+        acceleration_loss = pos_acceleration_loss + alpha_rot * rot_acceleration_loss
+        jerk_loss         = pos_jerk_loss + alpha_rot * rot_jerk_loss
+
+        smoothness_loss = (
+            weight_acceleration * acceleration_loss
+          + weight_jerk        * jerk_loss
+        ) * smoothness_multiplier
+
+
+        # Calculate the spline loss
+        l2_multiplier = 10.0
+        spline_loss = l2_multiplier * (l2_loss + smoothness_loss) + geodesic_loss + quaternion_dist
+
+        dtw_loss, hausdorff_loss, frechet_loss, chamfer_loss = None, None, None, None
+
+        # Uncomment these lines if you want to use the other losses
+        '''
+        dtw = DynamicTimeWarpingLoss()
+        dtw_loss, _ = dtw.forward(pred_trans_selected_values.reshape(-1, 3), targ_trans_selected_values.reshape(-1, 3))
+
+        hausdorff = HausdorffDistanceLoss()
+        hausdorff_loss, _ = hausdorff.forward(pred_trans_selected_values.reshape(-1, 3), targ_trans_selected_values.reshape(-1, 3))
+
+        frec = FrechetDistanceLoss()
+        frechet_loss, _ = frec.forward(pred_trans_selected_values.reshape(-1, 3), targ_trans_selected_values.reshape(-1, 3))
+
+        chamfer = ChamferDistanceLoss()
+        chamfer_loss, _ = chamfer.forward(pred_trans_selected_values.reshape(-1, 3), targ_trans_selected_values.reshape(-1, 3))
+        '''   
+
+        return spline_loss, l2_multiplier * l2_loss, l2_multiplier * smoothness_loss, geodesic_loss, quaternion_dist, dtw_loss, hausdorff_loss, frechet_loss, chamfer_loss
+
+
+class DynamicTimeWarpingLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def _dtw_distance(self, seq1: torch.Tensor, seq2: torch.Tensor) -> torch.Tensor:
+        """
+        Computes the DTW distance between two 2D tensors (T x D), 
+        where T is sequence length and D is feature dimension.
+        """
+        # seq1, seq2 shapes: (time_steps, feature_dim)
+        n, m = seq1.size(0), seq2.size(0)
+
+        # Cost matrix (pairwise distances between all elements)
+        cost = torch.zeros(n, m, device=seq1.device, dtype=seq1.dtype)
+        for i in range(n):
+            for j in range(m):
+                cost[i, j] = torch.norm(seq1[i] - seq2[j], p=2)
+
+        # Accumulated cost matrix
+        dist = torch.full((n + 1, m + 1), float('inf'), 
+                          device=seq1.device, dtype=seq1.dtype)
+        dist[0, 0] = 0.0
+
+        # Populate the DP table
+        for i in range(1, n + 1):
+            for j in range(1, m + 1):
+                dist[i, j] = cost[i - 1, j - 1] + torch.min(
+                    torch.min(
+                    dist[i - 1, j],   # Insertion
+                    dist[i, j - 1],   # Deletion
+                    ),
+                    dist[i - 1, j - 1]# Match
+                )
+
+        return dist[n, m]
+
+    def _loss(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the average DTW loss over a batch of sequences.
+        
+        pred, targ shapes: (batch_size, T, D)
+        """
+        # Ensure shapes match in batch dimension
+        assert pred.size(0) == targ.size(0), "Batch sizes must match."
+
+        # Compute DTW distance per sample in the batch
+        distances = []
+        for b in range(pred.size(0)):
+            seq1 = pred[b]
+            seq2 = targ[b]
+            dtw_val = self._dtw_distance(seq1, seq2)
+            distances.append(dtw_val)
+
+        # Stack and take mean to get scalar loss
+        dtw_loss = torch.stack(distances).mean()
+        return dtw_loss
+
+    def forward(self, pred: torch.Tensor, targ: torch.Tensor):
+        """
+        Returns a tuple: (loss, info_dict), 
+        where loss is a scalar tensor and info_dict is a dictionary 
+        of extra information (e.g., loss components).
+        """
+        loss = self._loss(pred, targ)
+
+        info = {
+            'dtw': loss.item()
+        }
+
+        return loss, info
+
+class HausdorffDistanceLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def _hausdorff_distance(self, set1: torch.Tensor, set2: torch.Tensor) -> torch.Tensor:
+        """
+        Computes the Hausdorff distance between two 2D tensors (N x D),
+        where N is the number of points and D is the feature dimension.
+
+        The Hausdorff distance H(A,B) between two sets A and B is defined as:
+            H(A, B) = max( h(A, B), h(B, A) ),
+        where
+            h(A, B) = max_{a in A} min_{b in B} d(a, b).
+
+        Here, d(a, b) is the Euclidean distance between points a and b.
+        """
+        # set1, set2 shapes: (num_points, feature_dim)
+        n, m = set1.size(0), set2.size(0)
+
+        # Compute pairwise distances
+        cost = torch.zeros(n, m, device=set1.device, dtype=set1.dtype)
+        for i in range(n):
+            for j in range(m):
+                cost[i, j] = torch.norm(set1[i] - set2[j], p=2)
+
+        # Forward direction: for each point in set1, find distance to closest point in set2
+        forward_min = cost.min(dim=1)[0]  # Shape (n,)
+        forward_hausdorff = forward_min.max()  # max over n
+
+        # Backward direction: for each point in set2, find distance to closest point in set1
+        backward_min = cost.min(dim=0)[0]  # Shape (m,)
+        backward_hausdorff = backward_min.max()  # max over m
+
+        # Hausdorff distance is the max of the two
+        hausdorff_dist = torch.max(forward_hausdorff, backward_hausdorff)
+        return hausdorff_dist
+
+    def _loss(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the average Hausdorff distance over a batch of point sets.
+
+        pred, targ shapes: (batch_size, N, D)
+        """
+        # Ensure shapes match in batch dimension
+        assert pred.size(0) == targ.size(0), "Batch sizes must match."
+
+        distances = []
+        for b in range(pred.size(0)):
+            set1 = pred[b]
+            set2 = targ[b]
+            h_dist = self._hausdorff_distance(set1, set2)
+            distances.append(h_dist)
+
+        # Stack and take mean to get scalar loss
+        hausdorff_loss = torch.stack(distances).mean()
+        return hausdorff_loss
+
+    def forward(self, pred: torch.Tensor, targ: torch.Tensor):
+        """
+        Returns a tuple: (loss, info_dict), 
+        where loss is a scalar tensor and info_dict is a dictionary 
+        of extra information (e.g., distance components).
+        """
+        loss = self._loss(pred, targ)
+
+        info = {
+            'hausdorff': loss.item()
+        }
+
+        return loss, info
+
+class FrechetDistanceLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def _frechet_distance(self, seq1: torch.Tensor, seq2: torch.Tensor) -> torch.Tensor:
+        """
+        Computes the (discrete) Fréchet distance between two 2D tensors (T x D),
+        where T is the sequence length and D is the feature dimension.
+
+        The Fréchet distance between two curves in discrete form can be computed
+        by filling in a DP table “ca” where:
+
+            ca[i, j] = max( d(seq1[i], seq2[j]),
+                            min(ca[i-1, j], ca[i, j-1], ca[i-1, j-1]) )
+
+        with boundary conditions handled appropriately.
+        Here, d(seq1[i], seq2[j]) is the Euclidean distance.
+        """
+        n, m = seq1.size(0), seq2.size(0)
+
+        # Cost matrix (pairwise distances between all elements)
+        cost = torch.zeros(n, m, device=seq1.device, dtype=seq1.dtype)
+        for i in range(n):
+            for j in range(m):
+                cost[i, j] = torch.norm(seq1[i] - seq2[j], p=2)
+
+        # DP matrix for the Fréchet distance
+        ca = torch.full((n, m), float('inf'), device=seq1.device, dtype=seq1.dtype)
+        ca[0, 0] = cost[0, 0]
+
+        # Initialize first row
+        for i in range(1, n):
+            ca[i, 0] = torch.max(ca[i - 1, 0], cost[i, 0])
+
+        # Initialize first column
+        for j in range(1, m):
+            ca[0, j] = torch.max(ca[0, j - 1], cost[0, j])
+
+        # Populate the DP table
+        for i in range(1, n):
+            for j in range(1, m):
+                ca[i, j] = torch.max(
+                    cost[i, j],
+                    torch.min(
+                        torch.min(
+                        ca[i - 1, j],
+                        ca[i, j - 1],
+                        ),
+                        ca[i - 1, j - 1]
+                    )
+                )
+
+        return ca[n - 1, m - 1]
+
+    def _loss(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the average Fréchet distance over a batch of sequences.
+
+        pred, targ shapes: (batch_size, T, D)
+        """
+        # Ensure shapes match in batch dimension
+        assert pred.size(0) == targ.size(0), "Batch sizes must match."
+
+        distances = []
+        for b in range(pred.size(0)):
+            seq1 = pred[b]
+            seq2 = targ[b]
+            fd_val = self._frechet_distance(seq1, seq2)
+            distances.append(fd_val)
+
+        # Stack and take mean to get scalar loss
+        frechet_loss = torch.stack(distances).mean()
+        return frechet_loss
+
+    def forward(self, pred: torch.Tensor, targ: torch.Tensor):
+        """
+        Returns a tuple: (loss, info_dict),
+        where loss is a scalar tensor and info_dict is a dictionary
+        of extra information (e.g., distance components).
+        """
+        loss = self._loss(pred, targ)
+        info = {
+            'frechet': loss.item()
+        }
+        return loss, info
+
+class ChamferDistanceLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def _chamfer_distance(self, set1: torch.Tensor, set2: torch.Tensor) -> torch.Tensor:
+        """
+        Computes the symmetrical Chamfer distance between
+        two 2D tensors (N x D), where N is the number of points
+        and D is the feature dimension.
+
+        The Chamfer distance between two point sets A and B is often defined as:
+
+            d_chamfer(A, B) = 1/|A| ∑_{a ∈ A} min_{b ∈ B} ‖a - b‖₂
+                              + 1/|B| ∑_{b ∈ B} min_{a ∈ A} ‖b - a‖₂,
+
+        where ‖·‖₂ is the Euclidean distance.
+        """
+        # set1, set2 shapes: (num_points, feature_dim)
+        n, m = set1.size(0), set2.size(0)
+
+        cost = torch.zeros(n, m, device=set1.device, dtype=set1.dtype)
+        for i in range(n):
+            for j in range(m):
+                cost[i, j] = torch.norm(set1[i] - set2[j], p=2)
+
+        # For each point in set1, find distance to the closest point in set2
+        forward_min = cost.min(dim=1)[0]   # shape: (n,)
+        forward_mean = forward_min.mean()
+
+        # For each point in set2, find distance to the closest point in set1
+        backward_min = cost.min(dim=0)[0]  # shape: (m,)
+        backward_mean = backward_min.mean()
+
+        chamfer_dist = forward_mean + backward_mean
+        return chamfer_dist
+
+    def _loss(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the average Chamfer distance over a batch of point sets.
+
+        pred, targ shapes: (batch_size, N, D)
+        """
+        # Ensure shapes match in batch dimension
+        assert pred.size(0) == targ.size(0), "Batch sizes must match."
+
+        distances = []
+        for b in range(pred.size(0)):
+            set1 = pred[b]
+            set2 = targ[b]
+            distance_val = self._chamfer_distance(set1, set2)
+            distances.append(distance_val)
+
+        # Combine into a single scalar
+        chamfer_loss = torch.stack(distances).mean()
+        return chamfer_loss
+
+    def forward(self, pred: torch.Tensor, targ: torch.Tensor):
+        """
+        Returns a tuple: (loss, info_dict),
+        where 'loss' is a scalar tensor and 'info_dict' is a dictionary
+        of extra information (e.g., distance components).
+        """
+        loss = self._loss(pred, targ)
+        info = {
+            'chamfer': loss.item()
+        }
+        return loss, info
+        
+
+def slerp(q1, q2, t):
+    """Spherical linear interpolation between two quaternions."""
+    q1 = q1 / np.linalg.norm(q1)
+    q2 = q2 / np.linalg.norm(q2)
+    dot = np.dot(q1, q2)
+
+    if dot < 0.0:
+        q2 = -q2
+        dot = -dot
+    # If dot is very close to 1, use linear interpolation
+    
+    if dot > 0.9995:
+        result = q1 + t * (q2 - q1)
+        result = result / np.linalg.norm(result)
+        return result
+    
+    theta_0 = np.arccos(dot)
+    theta = theta_0 * t
+    
+    q3 = q2 - q1 * dot
+    q3 = q3 / np.linalg.norm(q3)
+    return q1 * np.cos(theta) + q3 * np.sin(theta)
+
+def catmull_rom_spline_with_rotation(control_points, timepoints, horizon):
+    """Compute Catmull-Rom spline for both position and quaternion rotation."""
+    spline_points = []
+    # Extrapolate the initial points
+    if timepoints[0] != 0:
+        for t in range(timepoints[0]):
+            x = control_points[0][0]
+            y = control_points[0][1]
+            z = control_points[0][2]
+            q = control_points[0][3:7]
+            spline_points.append(np.concatenate([np.array([x, y, z]), q]))
+
+    #Linear interpolate between 0th and 1th control points
+    for t in np.linspace(0, 1, timepoints[1] - timepoints[0] + 1):
+        x = control_points[0][0] + t * (control_points[1][0] - control_points[0][0])
+        y = control_points[0][1] + t * (control_points[1][1] - control_points[0][1])
+        z = control_points[0][2] + t * (control_points[1][2] - control_points[0][2])
+        q = slerp(control_points[0][3:7], control_points[1][3:7], t)
+        spline_points.append(np.concatenate([np.array([x, y, z]), q]))
+
+
+    # Iterate over the control points
+    for i in range(1, len(control_points) - 2):
+        P0 = control_points[i-1][:3] 
+        P1 = control_points[i][:3]    
+        P2 = control_points[i+1][:3]  
+        P3 = control_points[i+2][:3]  
+        Q0 = control_points[i-1][3:7]  
+        Q1 = control_points[i][3:7]    
+        Q2 = control_points[i+1][3:7]  
+        Q3 = control_points[i+2][3:7]
+    
+        # Interpolate position (using Catmull-Rom spline)
+        for idx, t in enumerate(np.linspace(0, 1, timepoints[i+1] - timepoints[i] + 1)):
+            if idx == 0:
+                continue
+
+            x = 0.5 * ((2 * P1[0]) + (-P0[0] + P2[0]) * t + 
+                       (2 * P0[0] - 5 * P1[0] + 4 * P2[0] - P3[0]) * t**2 + 
+                       (-P0[0] + 3 * P1[0] - 3 * P2[0] + P3[0]) * t**3)
+            y = 0.5 * ((2 * P1[1]) + (-P0[1] + P2[1]) * t + 
+                       (2 * P0[1] - 5 * P1[1] + 4 * P2[1] - P3[1]) * t**2 + 
+                       (-P0[1] + 3 * P1[1] - 3 * P2[1] + P3[1]) * t**3)
+            z = 0.5 * ((2 * P1[2]) + (-P0[2] + P2[2]) * t + 
+                       (2 * P0[2] - 5 * P1[2] + 4 * P2[2] - P3[2]) * t**2 + 
+                       (-P0[2] + 3 * P1[2] - 3 * P2[2] + P3[2]) * t**3)
+            q = slerp(Q1, Q2, t)
+            spline_points.append(np.concatenate([np.array([x, y, z]), q]))
+        
+    #Linear interpolate between 2nd last and last control points
+    for idx, t in enumerate(np.linspace(0, 1, timepoints[-1] - timepoints[-2] + 1)):
+        if idx == 0:
+            continue
+        x = control_points[-2][0] + t * (control_points[-1][0] - control_points[-2][0])
+        y = control_points[-2][1] + t * (control_points[-1][1] - control_points[-2][1])
+        z = control_points[-2][2] + t * (control_points[-1][2] - control_points[-2][2])
+        q = slerp(control_points[-2][3:7], control_points[-1][3:7], t)
+        spline_points.append(np.concatenate([np.array([x, y, z]), q]))
+
+    # Extrapolate the rest of the points
+    if timepoints[-1] != horizon:
+        for t in range(timepoints[-1] + 1, horizon):
+            x = control_points[-1][0]
+            y = control_points[-1][1]
+            z = control_points[-1][2]
+            q = control_points[-1][3:7]
+            spline_points.append(np.concatenate([np.array([x, y, z]), q]))
+    
+    stacked_spline_points = np.stack(spline_points, axis=0)
+
+    if control_points.shape[1] != 7:
+        stacked_spline_points = np.concatenate([stacked_spline_points, np.zeros((stacked_spline_points.shape[0], 1))], axis=1)
+    
+    
+    return stacked_spline_points
+
+def catmull_rom_loss(trajectories, conditions, loss_fc):
+    '''
+        loss for catmull-rom interpolation
+    '''        
+    batch_size, horizon, transition = trajectories.shape
+    
+    # Extract known indices and values
+    known_indices = np.array(list(conditions.keys()), dtype=int)
+    
+    # candidate_no x batch_size x dim
+    known_values = np.stack([c.cpu().numpy() for c in conditions.values()], axis=0)
+    known_values = np.moveaxis(known_values, 0, 1)
+    
+    # Sort the timepoints
+    sorted_indices = np.argsort(known_indices)
+    known_indices = known_indices[sorted_indices]
+    known_values = known_values[:, sorted_indices]
+    spline_points = np.array([catmull_rom_spline_with_rotation(known_values[b], known_indices, horizon) for b in range(batch_size)])
+    
+    # Convert to tensor and move to the same device as trajectories
+    spline_points = torch.tensor(spline_points, dtype=torch.float64, device=trajectories.device)
+    assert spline_points.shape == trajectories.shape, f"Shape mismatch: {spline_points.shape} != {trajectories.shape}"
+    return loss_fc(spline_points, trajectories)
+
 Losses = {
     'l1': WeightedL1,
     'l2': WeightedL2,
     'value_l1': ValueL1,
     'value_l2': ValueL2,
+    'geodesic_l2': GeodesicL2Loss,
+    'rotation_translation': RotationTranslationLoss,
+    'spline': SplineLoss,
 }
diff --git a/diffuser/models/temporal.py b/diffuser/models/temporal.py
index e0b9e5c..0f7854a 100644
--- a/diffuser/models/temporal.py
+++ b/diffuser/models/temporal.py
@@ -17,18 +17,18 @@ class ResidualTemporalBlock(nn.Module):
         super().__init__()
 
         self.blocks = nn.ModuleList([
-            Conv1dBlock(inp_channels, out_channels, kernel_size),
-            Conv1dBlock(out_channels, out_channels, kernel_size),
+            Conv1dBlock(inp_channels, out_channels, kernel_size).to(dtype=torch.float64),
+            Conv1dBlock(out_channels, out_channels, kernel_size).to(dtype=torch.float64),
         ])
 
         self.time_mlp = nn.Sequential(
             nn.Mish(),
-            nn.Linear(embed_dim, out_channels),
+            nn.Linear(embed_dim, out_channels).to(dtype=torch.float64),
             Rearrange('batch t -> batch t 1'),
-        )
+        ).to(dtype=torch.float64)
 
-        self.residual_conv = nn.Conv1d(inp_channels, out_channels, 1) \
-            if inp_channels != out_channels else nn.Identity()
+        self.residual_conv = nn.Conv1d(inp_channels, out_channels, 1).to(dtype=torch.float64) \
+            if inp_channels != out_channels else nn.Identity().to(dtype=torch.float64)
 
     def forward(self, x, t):
         '''
@@ -37,7 +37,8 @@ class ResidualTemporalBlock(nn.Module):
             returns:
             out : [ batch_size x out_channels x horizon ]
         '''
-        out = self.blocks[0](x) + self.time_mlp(t)
+        
+        out = self.blocks[0](x) + self.time_mlp(t.double())
         out = self.blocks[1](out)
         return out + self.residual_conv(x)
 
@@ -49,11 +50,11 @@ class TemporalUnet(nn.Module):
         transition_dim,
         cond_dim,
         dim=32,
-        dim_mults=(1, 2, 4, 8),
+        dim_mults=(1, 2, 4),
     ):
         super().__init__()
 
-        dims = [transition_dim, *map(lambda m: dim * m, dim_mults)]
+        dims = [(transition_dim + cond_dim), *map(lambda m: dim * m, dim_mults)]
         in_out = list(zip(dims[:-1], dims[1:]))
         print(f'[ models/temporal ] Channel dimensions: {in_out}')
 
@@ -100,7 +101,7 @@ class TemporalUnet(nn.Module):
 
         self.final_conv = nn.Sequential(
             Conv1dBlock(dim, dim, kernel_size=5),
-            nn.Conv1d(dim, transition_dim, 1),
+            nn.Conv1d(dim, transition_dim, 1).to(dtype=torch.float64),
         )
 
     def forward(self, x, cond, time):
@@ -129,7 +130,6 @@ class TemporalUnet(nn.Module):
             x = upsample(x)
 
         x = self.final_conv(x)
-
         x = einops.rearrange(x, 'b t h -> b h t')
         return x
 
diff --git a/diffuser/utils/arrays.py b/diffuser/utils/arrays.py
index c3a9d24..96a7093 100644
--- a/diffuser/utils/arrays.py
+++ b/diffuser/utils/arrays.py
@@ -54,7 +54,7 @@ def batchify(batch):
 			1) converting np arrays to torch tensors and
 			2) and ensuring that everything has a batch dimension
 	'''
-	fn = lambda x: to_torch(x[None])
+	fn = lambda x: to_torch(x[None], dtype=torch.float64)
 
 	batched_vals = []
 	for field in batch._fields:
diff --git a/diffuser/utils/serialization.py b/diffuser/utils/serialization.py
index 6cc9db9..039eb64 100644
--- a/diffuser/utils/serialization.py
+++ b/diffuser/utils/serialization.py
@@ -19,7 +19,7 @@ def mkdir(savepath):
         return False
 
 def get_latest_epoch(loadpath):
-    states = glob.glob1(os.path.join(*loadpath), 'state_*')
+    states = glob.glob1(os.path.join(loadpath), 'state_*')
     latest_epoch = -1
     for state in states:
         epoch = int(state.replace('state_', '').replace('.pt', ''))
diff --git a/diffuser/utils/training.py b/diffuser/utils/training.py
index be3556e..c21e0f0 100644
--- a/diffuser/utils/training.py
+++ b/diffuser/utils/training.py
@@ -4,16 +4,24 @@ import numpy as np
 import torch
 import einops
 import pdb
+from tqdm import tqdm
+import wandb
+from pytorch3d.transforms import axis_angle_to_quaternion
 
 from .arrays import batch_to_device, to_np, to_device, apply_dict
 from .timer import Timer
 from .cloud import sync_logs
+from ..models.helpers import catmull_rom_spline_with_rotation
 
 def cycle(dl):
     while True:
         for data in dl:
             yield data
 
+def assert_no_nan_weights(model):
+    for name, param in model.named_parameters():
+        assert not torch.isnan(param).any(), f"NaN detected in parameter: {name}"
+
 class EMA():
     '''
         empirical moving average
@@ -71,13 +79,35 @@ class Trainer(object):
         self.gradient_accumulate_every = gradient_accumulate_every
 
         self.dataset = dataset
-        self.dataloader = cycle(torch.utils.data.DataLoader(
-            self.dataset, batch_size=train_batch_size, num_workers=1, shuffle=True, pin_memory=True
+        dataset_size = len(self.dataset)
+
+        # Read the indices from the .txt file
+        with open(os.path.join(results_folder, 'train_indices.txt'), 'r') as f:
+            self.train_indices = f.read()
+            self.train_indices = [int(i) for i in self.train_indices.split('\n') if i]
+        
+        with open(os.path.join(results_folder, 'val_indices.txt'), 'r') as f:
+            self.val_indices = f.read()
+            self.val_indices = [int(i) for i in self.val_indices.split('\n') if i]
+
+
+        self.train_dataset = torch.utils.data.Subset(self.dataset, self.train_indices)
+        self.val_dataset = torch.utils.data.Subset(self.dataset, self.val_indices)
+        self.train_dataloader = cycle(torch.utils.data.DataLoader(
+            self.train_dataset, batch_size=train_batch_size, num_workers=1, pin_memory=True, shuffle=False
+        ))
+
+        self.val_dataloader = cycle(torch.utils.data.DataLoader(
+            self.val_dataset, batch_size=train_batch_size, num_workers=1, pin_memory=True, shuffle=False
         ))
+
         self.dataloader_vis = cycle(torch.utils.data.DataLoader(
             self.dataset, batch_size=1, num_workers=0, shuffle=True, pin_memory=True
         ))
         self.renderer = renderer
+
+        
+
         self.optimizer = torch.optim.Adam(diffusion_model.parameters(), lr=train_lr)
 
         self.logdir = results_folder
@@ -88,6 +118,8 @@ class Trainer(object):
         self.reset_parameters()
         self.step = 0
 
+        self.log_to_wandb = False
+
     def reset_parameters(self):
         self.ema_model.load_state_dict(self.model.state_dict())
 
@@ -102,36 +134,129 @@ class Trainer(object):
     #-----------------------------------------------------------------------------#
 
     def train(self, n_train_steps):
-
+        # Save the indices as .txt files
+        with open(os.path.join(self.logdir, 'train_indices.txt'), 'w') as f:
+            for idx in self.train_indices:
+                f.write(f"{idx}\n")
+        with open(os.path.join(self.logdir, 'val_indices.txt'), 'w') as f:
+            for idx in self.val_indices:
+                f.write(f"{idx}\n")
+        
         timer = Timer()
-        for step in range(n_train_steps):
+        torch.autograd.set_detect_anomaly(True)
+
+        # Setup wandb
+        if self.log_to_wandb:
+            wandb.init(
+                project='trajectory-generation', 
+                config={'lr': self.optimizer.param_groups[0]['lr'], 'batch_size': self.batch_size, 'gradient_accumulate_every': self.gradient_accumulate_every},
+            )
+
+        for step in tqdm(range(n_train_steps)):
+            
+            mean_train_loss = 0.0
             for i in range(self.gradient_accumulate_every):
-                batch = next(self.dataloader)
+                batch = next(self.train_dataloader)
                 batch = batch_to_device(batch)
-
-                loss, infos = self.model.loss(*batch)
+                
+                loss, infos = self.model.loss(x=batch.trajectories, cond=batch.conditions)
                 loss = loss / self.gradient_accumulate_every
+                mean_train_loss += loss.item()
                 loss.backward()
 
+            if self.log_to_wandb:
+                wandb.log({
+                    'step': self.step,
+                    'train/loss': mean_train_loss
+                })
+
+            # torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+
             self.optimizer.step()
             self.optimizer.zero_grad()
 
+            assert_no_nan_weights(self.model)
+
             if self.step % self.update_ema_every == 0:
                 self.step_ema()
 
             if self.step % self.save_freq == 0:
-                label = self.step // self.label_freq * self.label_freq
+                label = self.step
+                print(f'Saving model at step {self.step}...')
                 self.save(label)
 
             if self.step % self.log_freq == 0:
-                infos_str = ' | '.join([f'{key}: {val:8.4f}' for key, val in infos.items()])
-                print(f'{self.step}: {loss:8.4f} | {infos_str} | t: {timer():8.4f}')
+                val_losses = []
+                lin_int_losses = []
+
+                val_infos_list = []
+                lin_int_infos_list = []
+
+                catmull_losses = []
+                catmull_infos_list = []
+
+                for _ in range(len(self.val_indices)):
+                    val_batch = next(self.val_dataloader)
+                    val_batch = batch_to_device(val_batch)
+
+                    traj = self.model.forward(val_batch.conditions, horizon=val_batch.trajectories.shape[1])
+                    val_loss, val_infos = self.model.loss_fn(traj, val_batch.trajectories, cond=val_batch.conditions)
+
+                    val_losses.append(val_loss.item())
+                    val_infos_list.append({key: val for key, val in val_infos.items()})
+                    
+
+                    (lin_int_loss, lin_int_infos), lin_int_traj = self.linear_interpolation_loss(
+                        val_batch.trajectories, val_batch.conditions, self.model.loss_fn
+                    )
+                    lin_int_losses.append(lin_int_loss.item())
+                    lin_int_infos_list.append({key: val for key, val in lin_int_infos.items()})
+
+                    (catmull_loss, catmull_infos), catmull_traj = self.catmull_rom_loss(
+                        val_batch.trajectories, val_batch.conditions, self.model.loss_fn
+                    )
+
+                    catmull_losses.append(catmull_loss.item())
+                    catmull_infos_list.append(catmull_infos)
+                
+                avg_val_loss = np.mean(val_losses)
+                avg_lin_int_loss = np.mean(lin_int_losses)
+
+                val_infos = {key: np.mean([info[key] for info in val_infos_list]) for key in val_infos_list[0].keys()}
+                lin_int_infos = {key: np.mean([info[key] for info in lin_int_infos_list]) for key in lin_int_infos_list[0].keys()}
 
-            if self.step == 0 and self.sample_freq:
-                self.render_reference(self.n_reference)
+                avg_catmull_loss = np.mean(catmull_losses)
+                catmull_infos = {key: np.mean([info[key] for info in catmull_infos_list]) for key in catmull_infos_list[0].keys()}
 
-            if self.sample_freq and self.step % self.sample_freq == 0:
-                self.render_samples(n_samples=self.n_samples)
+                val_infos_str = ' | '.join([f'{key}: {val:8.4f}' for key, val in val_infos.items()])
+                lin_int_infos_str = ' | '.join([f'{key}: {val:8.4f}' for key, val in lin_int_infos.items()])
+                catmull_infos_str = ' | '.join([f'{key}: {val:8.4f}' for key, val in catmull_infos.items()])
+                
+
+                infos_str = ' | '.join([f'{key}: {val:8.4f}' for key, val in infos.items()])
+                print("Learning Rate: ", self.optimizer.param_groups[0]['lr'])
+                print(f'Step {self.step}: {loss * self.gradient_accumulate_every:8.4f} | {infos_str} | t: {timer():8.4f}')
+                print(f'Validation - {self.step}: {avg_val_loss:8.4f} | {val_infos_str} | t: {timer():8.4f}')
+                print(f'Linear Interpolation Loss - {self.step}: {avg_lin_int_loss:8.4f} | {lin_int_infos_str} | t: {timer():8.4f}')
+                print(f'Catmull Rom Loss - {self.step}: {avg_catmull_loss:8.4f} | {catmull_infos_str} | t: {timer():8.4f}')
+                print()
+
+                if self.log_to_wandb:
+                    wandb.log({
+                        'step': self.step,
+                        'val/loss': avg_val_loss,
+                        'val/linear_interp/loss': avg_lin_int_loss,
+                        'val/linear_interp/quaternion dist.': lin_int_infos['quat. dist.'],
+                        'val/linear_interp/euclidean dist.': lin_int_infos['trans. error'],
+                        'val/linear_interp/geodesic loss': lin_int_infos['geodesic dist.'],
+                        'val/catmull_rom/loss': avg_catmull_loss,
+                        'val/catmull_rom/quaternion dist.': catmull_infos['quat. dist.'],
+                        'val/catmull_rom/euclidean dist.': catmull_infos['trans. error'],
+                        'val/catmull_rom/geodesic loss': catmull_infos['geodesic dist.'],
+                        'val/quaternion dist.': val_infos['quat. dist.'],
+                        'val/euclidean dist.': val_infos['trans. error'],
+                        'val/geodesic loss': val_infos['geodesic dist.'],
+                    })
 
             self.step += 1
 
@@ -186,15 +311,6 @@ class Trainer(object):
         normed_observations = trajectories[:, :, self.dataset.action_dim:]
         observations = self.dataset.normalizer.unnormalize(normed_observations, 'observations')
 
-        # from diffusion.datasets.preprocessing import blocks_cumsum_quat
-        # # observations = conditions + blocks_cumsum_quat(deltas)
-        # observations = conditions + deltas.cumsum(axis=1)
-
-        #### @TODO: remove block-stacking specific stuff
-        # from diffusion.datasets.preprocessing import blocks_euler_to_quat, blocks_add_kuka
-        # observations = blocks_add_kuka(observations)
-        ####
-
         savepath = os.path.join(self.logdir, f'_sample-reference.png')
         self.renderer.composite(savepath, observations)
 
@@ -225,9 +341,6 @@ class Trainer(object):
             # [ 1 x 1 x observation_dim ]
             normed_conditions = to_np(batch.conditions[0])[:,None]
 
-            # from diffusion.datasets.preprocessing import blocks_cumsum_quat
-            # observations = conditions + blocks_cumsum_quat(deltas)
-            # observations = conditions + deltas.cumsum(axis=1)
 
             ## [ n_samples x (horizon + 1) x observation_dim ]
             normed_observations = np.concatenate([
@@ -238,10 +351,70 @@ class Trainer(object):
             ## [ n_samples x (horizon + 1) x observation_dim ]
             observations = self.dataset.normalizer.unnormalize(normed_observations, 'observations')
 
-            #### @TODO: remove block-stacking specific stuff
-            # from diffusion.datasets.preprocessing import blocks_euler_to_quat, blocks_add_kuka
-            # observations = blocks_add_kuka(observations)
-            ####
-
             savepath = os.path.join(self.logdir, f'sample-{self.step}-{i}.png')
             self.renderer.composite(savepath, observations)
+
+    def linear_interpolation_loss(self, trajectories, conditions, loss_fc, scene_id=None, norm_params=None):
+        batch_size, horizon, transition = trajectories.shape
+
+        # Extract known indices and values
+        known_indices = np.array(list(conditions.keys()), dtype=int)
+        # candidate_no x batch_size x dim
+        known_values = np.stack([c.cpu().numpy() for c in conditions.values()], axis=0)
+        known_values = np.moveaxis(known_values, 0, 1)
+
+        # Create time steps for interpolation
+        time_steps = np.linspace(0, horizon, num=horizon)
+
+        # Perform interpolation across all dimensions at once
+        linear_int_arr = np.array([[
+            np.interp(time_steps, known_indices, known_values[b, :, dim])
+            for dim in range(transition)] 
+            for b in range(batch_size)]
+        ).T  # Transpose to match shape (horizon, transition)
+
+        # Convert to tensor and move to the same device as trajectories
+        linear_int_arr = np.transpose(linear_int_arr, axes=[2, 0, 1])
+        linear_int_tensor = torch.tensor(linear_int_arr, dtype=torch.float64, device=trajectories.device)
+        
+        return loss_fc(linear_int_tensor, trajectories, cond=conditions, scene_id=scene_id, norm_params=norm_params), linear_int_tensor       
+
+
+    def catmull_rom_loss(self, trajectories, conditions, loss_fc, scene_id=None, norm_params=None):
+        '''
+            loss for catmull-rom interpolation
+        '''        
+    
+        batch_size, horizon, transition = trajectories.shape
+
+        # Extract known indices and values
+        known_indices = np.array(list(conditions.keys()), dtype=int)
+        # candidate_no x batch_size x dim
+        known_values = np.stack([c.cpu().numpy() for c in conditions.values()], axis=0)
+        known_values = np.moveaxis(known_values, 0, 1)
+
+        # Sort the timepoints
+        sorted_indices = np.argsort(known_indices)
+        known_indices = known_indices[sorted_indices]
+        known_values = known_values[:, sorted_indices]
+
+        spline_points = np.array([catmull_rom_spline_with_rotation(known_values[b], known_indices, horizon) for b in range(batch_size)])
+        
+        # Convert to tensor and move to the same device as trajectories
+        spline_points = torch.tensor(spline_points, dtype=torch.float64, device=trajectories.device)
+
+        assert spline_points.shape == trajectories.shape, f"Shape mismatch: {spline_points.shape} != {trajectories.shape}"
+
+        return loss_fc(spline_points, trajectories, cond=conditions, scene_id=scene_id, norm_params=norm_params), spline_points
+        
+
+
+
+
+        
+
+
+
+          
+        
+          
diff --git a/scripts/train.py b/scripts/train.py
index 2c5f299..6728d6f 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -108,6 +108,7 @@ utils.report_parameters(model)
 
 print('Testing forward...', end=' ', flush=True)
 batch = utils.batchify(dataset[0])
+
 loss, _ = diffusion.loss(*batch)
 loss.backward()
 print('✓')