intial commint

basis_gn.py

@@ -0,0 +1,349 @@
+"""
+@brief:     Basis generators in PyTorch
+"""
+from typing import Tuple
+
+import torch
+
+
+class UniBSplineBasis(torch.nn.Module):
+
+    def __init__(self,
+                 num_basis: int = 10,
+                 degree_p: int = 3,
+                 dtype: torch.dtype = torch.float32,
+                 device: torch.device = 'cpu',
+                 **kwargs):
+        """
+        Constructor for basis class
+        Args:
+            num_basis: number of basis functions
+            dtype: torch data type
+            device: torch device to run on
+        """
+        super().__init__()
+
+        # Internal number of basis
+        self.num_basis = num_basis
+
+        self.degree_p = degree_p
+        self.init_cond_order = kwargs.get("init_condition_order", 0)
+        self.end_cond_order = kwargs.get("end_condition_order", 0)
+
+        self.num_ctrlp = num_basis + self.init_cond_order + self.end_cond_order
+        # number of knots needed, with respect to B-sp degree and number of
+        # control points ( num_basis + init_cond_order+end_cond_order)
+        num_knots = self.degree_p + 1 + self.num_ctrlp
+        num_knots_non_rep_inside_1 = num_knots - 2 * self.degree_p
+        # uniform knots vector
+        knots_vec = torch.linspace(0, 1, num_knots_non_rep_inside_1,
+                                        dtype=dtype, device=device)
+        knots_prev = torch.zeros(self.degree_p, dtype=dtype, device=device)
+        knots_pro = torch.ones(self.degree_p, dtype=dtype, device=device)
+        knots_vec = torch.cat([knots_prev, knots_vec, knots_pro])
+        self.register_buffer("knots_vec", knots_vec, persistent=False)
+
+        tau = kwargs.get("tau")
+        self.register_buffer('tau',
+                             torch.tensor(tau, dtype=dtype, device=device),
+                             persistent=False)
+
+    @property
+    def device(self):
+        return self.knots_vec.device
+
+    @property
+    def dtype(self):
+        return self.knots_vec.dtype
+
+    def time2phase(self, times: torch.Tensor) -> torch.Tensor:
+        """
+        scaling time into [0,1] range phase
+        :param times:
+        :return:
+        """
+        # Shape of times:
+        # [*add_dim, num_times]
+
+        # tau = times[..., -1]
+        tau = times.reshape(-1)[-1]
+        self.tau.copy_(tau)
+        phase = torch.clip(times / self.tau[..., None], 0, 1)
+        return phase
+
+    def basis(self, times: torch.Tensor) -> torch.Tensor:
+        """
+        compute evaluated b-spline basis at given time points
+        :param times:
+        :return:
+        """
+
+        # Shape of times:
+        # [*add_dim, num_times]
+        #
+        # Shape of basis:
+        # [*add_dim, num_times, num_ctrlp]
+
+        # phase = self.phase_generator.phase(times)
+        phase = self.time2phase(times)
+
+        basis = [self._basis_function(i, self.degree_p, self.knots_vec, phase)
+                 for i in range(self.num_ctrlp)]
+        basis = torch.stack(basis, dim=-1)
+
+        return basis
+
+    def _basis_function(self, i, k, knots, u, **kwargs):
+        """
+        recursive construct of B-spline basis using de Boor's algorithm
+
+        :param i: basis index
+        :param k: degree
+        :param u: evaluate time point
+        :param knots: knots vector
+        :return: vector of shape [num_eval_points]
+        """
+
+        if k == 0:
+            num_ctrlp = kwargs.get("num_ctrlp", self.num_ctrlp)
+            if i == num_ctrlp - 1:
+                # with regard to original definition, each span is defined as \
+                # left closed and right open interval [v_i, v_i+1), which makes\
+                # the value at right end always 0. It is undesired,so that we \
+                # need to handle the last basis specially
+                b0 = torch.where((u >= knots[i]) & (u <= knots[i + 1]), 1, 0)
+            else:
+                b0 = torch.where((u >= knots[i]) & (u < knots[i + 1]), 1, 0)
+            return torch.as_tensor(b0, dtype=self.dtype, device=self.device)
+        else:
+            denom1 = knots[i + k] - knots[i]
+            term1 = 0.0 if denom1 == 0 else (u - knots[i]) / denom1 * \
+                                            self._basis_function(i, k - 1,
+                                                                 knots, u,
+                                                                 **kwargs)
+            denom2 = knots[i + k + 1] - knots[i + 1]
+            term2 = 0.0 if denom2 == 0 else (knots[i + k + 1] - u) / denom2 * \
+                                            self._basis_function(i + 1, k - 1,
+                                                                 knots, u,
+                                                                 **kwargs)
+            return term1 + term2
+
+    def vel_basis(self, times: torch.Tensor) -> torch.Tensor:
+        """
+        Directly get the basis of velocity B-spline
+        :param times:
+        :return:
+        """
+
+        # phase = self.phase_generator.phase(times)
+        phase = self.time2phase(times)
+
+        # for clamped uni B-spline
+        vel_nots_vec = self.knots_vec[1:-1]
+        basis = \
+            [self._basis_function(i, self.degree_p - 1, vel_nots_vec, phase,
+                                  num_ctrlp=self.num_ctrlp - 1)
+             for i in range(self.num_ctrlp - 1)]
+        basis = torch.stack(basis, dim=-1)
+        if self.goal_basis:
+            gb = torch.ones_like(phase, dtype=self.dtype, device=self.device)[
+                ..., None]
+            basis = torch.cat([basis, gb], dim=-1)
+        return basis
+
+    def acc_basis(self, times: torch.Tensor) -> torch.Tensor:
+        """
+        Directly get the basis of acceleration B-spline
+        :param times:
+        :return:
+        """
+
+        # phase = self.phase_generator.phase(times)
+        phase = self.time2phase(times)
+
+        acc_knots_vec = self.knots_vec[2: -2]
+
+        basis = [
+            self._basis_function(i, self.degree_p - 2, acc_knots_vec, phase,
+                                 num_ctrlp=self.num_ctrlp - 2)
+            for i in range(self.num_ctrlp - 2)]
+        basis = torch.stack(basis, dim=-1)
+
+        return basis
+
+    def velocity_control_points(self, ctrl_pts: torch.Tensor):
+        """
+        given the position control points (parameter), return the velocity control
+        points for vel B-spline as linear combination of position control points.
+
+        :param ctrl_pts: vector of position control points
+        :return: velocity control points
+        """
+        # diff shape: [*add_dim, num_dof, num_ctrlp-1]
+        diff = ctrl_pts[..., 1:] - ctrl_pts[..., :-1]
+        # shape: [num_basis-1]
+        delta = self.knots_vec[
+                1 + self.degree_p: self.num_ctrlp + self.degree_p] - \
+                self.knots_vec[1: self.num_ctrlp]
+        diff = diff * (1 / delta)
+        return diff * self.degree_p
+
+    def acceleration_control_points(self, ctrl_pts: torch.Tensor):
+        """
+        given the position control points (parameter), return the acceleration
+        control points for acc B-spline as linear combination of position
+        control points.
+
+        :param ctrl_pts: vector of position control points
+        :return: velocity control points
+        """
+        # shape: [*add_dim, num_dof, num_ctrlp-1]
+        vel_ctrl_pts = self.velocity_control_points(ctrl_pts)
+        # shape: [*add_dim, num_dof, num_ctrlp-2]
+        diff = vel_ctrl_pts[..., 1:] - vel_ctrl_pts[..., :-1]
+        # shape: [num_ctrlp-2]
+        # delta = self.knots_vec[2+self.degree_p: self.num_ctrlp+self.degree_p-1]\
+        #     - self.knots_vec[2: self.num_ctrlp-1]
+        delta = self.knots_vec[
+                2 + self.degree_p: self.num_ctrlp + self.degree_p] \
+                - self.knots_vec[2: self.num_ctrlp]
+        diff = diff * (1 / delta)
+        return diff * (self.degree_p - 1)
+
+    def compute_init_params(self, init_pos, init_vel, **kwargs):
+        """
+        Given initial condition, compute corresponding the first control points
+        :param init_pos:
+        :param init_vel:
+        :param kwargs:
+        :return:
+        """
+
+        # Shape of init_pos:
+        # [*add_dim, num_dof]
+        #
+        # Shape of init_vel:
+        # [*add_dim, num_dof]
+        #
+        # return shape:
+        # [*add_dim, num_dof, init_cond_order]
+
+        if self.init_cond_order == 0:
+            return None
+
+        para_init_p = init_pos
+        para_init = para_init_p[..., None]
+
+        if self.init_cond_order == 2:
+            para_init_v = \
+                torch.einsum("...i,...->...i", init_vel,
+                             self.tau) * \
+                (self.knots_vec[1 + self.degree_p] - self.knots_vec[1]) \
+                / self.degree_p + para_init_p
+            para_init = torch.cat([para_init, para_init_v[..., None]], dim=-1)
+
+        return para_init
+
+    def compute_end_params(self, end_pos, end_vel, **kwargs):
+        """
+        Given end condition, compute corresponding the last control points
+        :param end_pos:
+        :param end_vel:
+        :param kwargs:
+        :return:
+        """
+        # Shape of end_pos:
+        # [*add_dim, num_dof]
+        #
+        # Shape of end_vel:
+        # [*add_dim, num_dof]
+        #
+        # return shape:
+        # [*add_dim, num_dof, init_cond_order]
+
+        if self.end_cond_order == 0:
+            return None
+
+        para_end_p = end_pos
+        para_end = para_end_p[..., None]
+
+        if self.end_cond_order == 2:
+            para_end_v = para_end_p - \
+                         torch.einsum("...i,...->...i", end_vel,
+                                      self.tau) * \
+                         (self.knots_vec[self.num_ctrlp - 1 + self.degree_p] -
+                          self.knots_vec[self.num_ctrlp - 1]) * self.degree_p
+            # para_end_v = para_end_p - (end_vel * self.phase_generator.tau) * \
+            #               (self.knots_vec[self.num_ctrlp - 1 + self.degree_p] -
+            #                self.knots_vec[self.num_ctrlp-1]) * self.degree_p
+            para_end = torch.cat([para_end_v[..., None], para_end], dim=-1)
+
+        return para_end
+
+    def basis_multi_dofs(self,
+                         times: torch.Tensor,
+                         num_dof: int) -> torch.Tensor:
+        """
+        Interface to generate value of single basis function at given time
+        points
+        Args:
+            times: times in Tensor
+            num_dof: num of Degree of freedoms
+        Returns:
+            basis_multi_dofs: Multiple DoFs basis functions in Tensor
+
+        """
+        # Shape of time
+        # [*add_dim, num_times]
+        #
+        # Shape of basis_multi_dofs
+        # [*add_dim, num_dof * num_times, num_dof * num_basis]
+
+        # Extract additional dimensions
+        add_dim = list(times.shape[:-1])
+
+        # Get single basis, shape: [*add_dim, num_times, num_ctrlp]
+        basis_single_dof = self.basis(times)
+        # num_times = basis_single_dof.shape[-2]
+        num_times = times.shape[-1]
+
+        # shape: [*add_dim, num_times, num_basis]
+        basis_single_dof_ = basis_single_dof[..., self.init_cond_order:
+                                                  self.num_ctrlp - self.end_cond_order]
+        # Multiple Dofs, shape:
+        # [*add_dim, num_dof * num_times, num_dof * num_basis]
+        basis_multi_dofs = torch.zeros(*add_dim, num_dof * num_times,
+                                       num_dof * self.num_basis,
+                                       dtype=self.dtype,
+                                       device=self.device)
+
+        # Assemble
+        for i in range(num_dof):
+            row_indices = slice(i * num_times, (i + 1) * num_times)
+            col_indices = slice(i * self.num_basis, (i + 1) * self.num_basis)
+            basis_multi_dofs[..., row_indices, col_indices] = basis_single_dof_
+
+        # Return
+        return basis_multi_dofs
+
+    def show_basis(self, plot=False) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute basis function values for debug usage
+        The times are in the range of [delay - tau, delay + 2 * tau]
+
+        Returns: basis function values
+
+        """
+        times = torch.linspace(0, 1, steps=1000)
+        basis_values = self.basis(times)
+        if plot:
+            import matplotlib.pyplot as plt
+            plt.figure()
+            for i in range(basis_values.shape[-1]):
+                plt.plot(times, basis_values[:, i], label=f"basis_{i}")
+            plt.grid()
+            plt.legend()
+            plt.axvline(x=0, linestyle='--', color='k', alpha=0.3)
+            plt.axvline(x=1, linestyle='--', color='k', alpha=0.3)
+            plt.show()
+        return times, basis_values

beast.py

@@ -0,0 +1,280 @@
+from .bspline_factory import SplineFactory
+import torch
+from addict import Dict
+from .utils import continuous_to_discrete, discrete_to_continuous, normalize_tensor, denormalize_tensor, tensor_linspace
+import numpy as np
+import matplotlib.pyplot as plt
+import einops
+
+from transformers.processing_utils import ProcessorMixin
+
+from functools import wraps
+
+def autocast_float32(fn):
+    @wraps(fn)
+    def wrapped(*args, **kwargs):
+        with torch.cuda.amp.autocast(dtype=torch.float32):
+            return fn(*args, **kwargs)
+    return wrapped
+
+class BeastTokenizer(torch.nn.Module, ProcessorMixin):
+    """
+    B-spline based tokenizer for trajectory encoding/decoding.
+    
+    Converts continuous trajectories to discrete tokens and vice versa using B-splines.
+    Supports continuous and discrete representations of trajectories.
+    Supports sperate handling for continous action and discrete state (e.g., binarized gripper state).
+    """
+    
+    # Class constants
+    DEFAULT_DT = 0.01  # 100 Hz sampling rate
+    
+    def __init__(self, num_dof=1, num_basis=10, seq_len=50, vocab_size=256,
+                 degree_p=4, gripper_zero_order=False, gripper_dof=1, init_cond_order=0, 
+                 end_cond_order=0, enforce_init_pos=True, device="cuda"):
+        super().__init__()
+        
+        # Store core parameters
+        self.device = device
+        self.seq_length = seq_len
+        self.vocab_size = vocab_size
+        self.num_basis = num_basis
+        self.enforce_init_pos = enforce_init_pos
+        self.init_cond_order = init_cond_order
+        self.end_cond_order = end_cond_order
+        self.dt = self.DEFAULT_DT
+        self.init_pos = None
+        
+        # Calculate DOF distribution
+        self.gripper_dof = gripper_dof if gripper_zero_order else 0
+        self.joint_dof = num_dof - self.gripper_dof
+        self.num_dof = self.joint_dof + self.gripper_dof
+        
+        # Initialize spline components
+        self.bsp = self._create_bsplines(self.joint_dof, degree_p)
+        self.gripper_bsp = self._create_bsplines(self.gripper_dof, 0) if gripper_zero_order else None
+        
+        # Setup time grid and weight bounds
+        # Working with normalized time [0, 1]
+        self.times = tensor_linspace(0, 1.0, seq_len).to(device)
+        self._initialize_weight_bounds()
+        
+        self.to(self.device)
+    
+    def _create_bsplines(self, num_dof, degree_p):
+        """Create motion primitive for joint trajectories."""
+        config = Dict({
+            'mp_type': 'uni_bspline',
+            'device': self.device,
+            'num_dof': num_dof,
+            'tau': 1.0,
+            'mp_args': {
+                'num_basis': self.num_basis,
+                'degree_p': degree_p,
+                'init_condition_order': self.init_cond_order,
+                'end_condition_order': self.end_cond_order,
+                'dt': self.dt
+            }
+        })
+        return SplineFactory.init_splines(**config)
+    
+    def _initialize_weight_bounds(self):
+        """Initialize weight bounds for normalization."""
+        total_params = self.num_dof * self.num_basis
+        self.register_buffer("w_min", -1.0 * torch.ones(total_params))
+        self.register_buffer("w_max", 1.0 * torch.ones(total_params))
+    
+    def _get_repeated_times(self, batch_size):
+        """Get time tensor repeated for batch processing."""
+        return einops.repeat(self.times, 't -> b t', b=batch_size)
+
+    @autocast_float32 
+    def _learn_trajectory_params(self, times, trajs):
+        """Learn B-spline parameters from trajectories."""
+        # Learn joint parameters
+        joint_params = self.bsp.learn_mp_params_from_trajs(times, trajs[..., :self.joint_dof])
+        
+        # Learn gripper parameters if applicable
+        if self.gripper_bsp is not None:
+            gripper_trajs = trajs[..., -self.gripper_dof:]
+            gripper_params = self.gripper_bsp.learn_mp_params_from_trajs(times, gripper_trajs)
+            joint_params['params'] = torch.cat([joint_params['params'], gripper_params['params']], dim=-1)
+        
+        return joint_params
+
+    @autocast_float32 
+    def _reconstruct_trajectory(self, params, times):
+        """Reconstruct trajectory from B-spline parameters."""
+        # Reconstruct joint trajectory
+        joint_params = params[..., :self.joint_dof * self.num_basis]
+        self.bsp.update_inputs(times=times, params=joint_params)
+        position = self.bsp.get_traj_pos()
+        
+        # Reconstruct gripper trajectory if applicable
+        if self.gripper_bsp is not None:
+            gripper_params = params[..., -self.gripper_dof * self.num_basis:]
+            self.gripper_bsp.update_inputs(times=times, params=gripper_params)
+            gripper_pos = self.gripper_bsp.get_traj_pos()
+            position = torch.cat([position, gripper_pos], dim=-1)
+        
+        return position
+    
+    def _apply_initial_position_constraint(self, params, init_pos):
+        """Apply initial position constraint to parameters."""
+        if not self.init_pos or init_pos is None:
+            return params
+            
+        # Reshape to access individual basis functions
+        reshaped_params = einops.rearrange(params, "b (d t) -> b t d", t=self.num_basis, d=self.num_dof)
+        
+        # Set initial position for joint DOFs
+        reshaped_params[:, 0, :self.joint_dof] = init_pos[:, :self.joint_dof]
+        
+        return einops.rearrange(reshaped_params, "b t d -> b (d t)")
+
+    @autocast_float32 
+    def compute_weights(self, demos):
+        """Compute B-spline weights from demonstration trajectories."""
+        times = self._get_repeated_times(demos.shape[0])
+        weights = self.bsp.learn_mp_params_from_trajs(times, demos)['params']
+        return weights
+    
+    def update_weights_bounds_per_batch(self, weights):
+        """Update weight bounds based on batch statistics."""
+        weights = weights.reshape(-1, self.num_dof * self.num_basis)
+        batch_min = weights.min(dim=0)[0]
+        batch_max = weights.max(dim=0)[0]
+        
+        # Update bounds with small tolerance
+        tolerance = 1e-4
+        smaller_mask = batch_min < (self.w_min - tolerance)
+        larger_mask = batch_max > (self.w_max + tolerance)
+        
+        if torch.any(smaller_mask):
+            self.w_min[smaller_mask] = batch_min[smaller_mask]
+        if torch.any(larger_mask):
+            self.w_max[larger_mask] = batch_max[larger_mask]
+    
+    def update_times(self, times):
+        """Update time grid."""
+        self.times = times
+    
+    @torch.no_grad() 
+    @autocast_float32
+    def encode_discrete(self, trajs, update_bounds=False, init_p=None):
+        """Encode trajectories to discrete tokens."""
+        times = self._get_repeated_times(trajs.shape[0])
+        params_dict = self._learn_trajectory_params(times, trajs)
+        
+        if update_bounds:
+            self.update_weights_bounds_per_batch(params_dict['params'])
+        
+        # Clamp parameters to bounds
+        params = torch.clamp(params_dict['params'], min=self.w_min, max=self.w_max)
+        
+        # Convert to discrete tokens
+        tokens = continuous_to_discrete(params, min_val=self.w_min, max_val=self.w_max, num_bins=self.vocab_size)
+        tokens = einops.rearrange(tokens, 'b (d t) -> b (t d)', t=self.num_basis, d=self.num_dof)
+        
+        return tokens
+    
+    @torch.no_grad()
+    @autocast_float32
+    def decode_discrete(self, tokens, times=None, init_pos=None):
+        """Decode discrete tokens to trajectories."""
+        # Reshape tokens and convert to continuous parameters
+        tokens = einops.rearrange(tokens, 'b (t d) -> b (d t)', t=self.num_basis, d=self.num_dof)
+        params = discrete_to_continuous(tokens, min_val=self.w_min, max_val=self.w_max, num_bins=self.vocab_size)
+        
+        if times is None:
+            times = self._get_repeated_times(params.shape[0])
+        
+        # Apply initial position constraint if specified
+        params = self._apply_initial_position_constraint(params, init_pos)
+        
+        return self._reconstruct_trajectory(params, times)
+    
+    @torch.no_grad() 
+    @autocast_float32
+    def encode_continuous(self, trajs, update_bounds=False):
+        """Encode trajectories to continuous tokens (normalized parameters)."""
+        times = self._get_repeated_times(trajs.shape[0])
+        params_dict = self._learn_trajectory_params(times, trajs)
+        
+        if update_bounds:
+            self.update_weights_bounds_per_batch(params_dict['params'])
+        
+        # Normalize parameters
+        tokens = normalize_tensor(params_dict['params'], w_min=self.w_min, w_max=self.w_max)
+        
+        return tokens
+    
+    @torch.no_grad()
+    @autocast_float32
+    def decode_continuous(self, params, times=None, init_pos=None):
+        """Decode continuous tokens (normalized parameters) to trajectories."""
+        # Denormalize parameters
+        params = denormalize_tensor(params, w_min=self.w_min, w_max=self.w_max)
+        
+        if times is None:
+            times = self._get_repeated_times(params.shape[0])
+        
+        # Apply initial position constraint if specified
+        params = self._apply_initial_position_constraint(params, init_pos)
+        
+        return self._reconstruct_trajectory(params, times)
+
+    @autocast_float32 
+    def compute_reconstruction_error(self, raw_traj):
+        """Compute reconstruction error for trajectory."""
+        if len(raw_traj.shape) == 2:
+            raw_traj = raw_traj.unsqueeze(-1)
+        
+        tokens, _ = self.encode_discrete(raw_traj)
+        reconstructed = self.decode_discrete(tokens)
+        error = torch.mean((raw_traj - reconstructed) ** 2)
+        
+        return error
+    
+    def _plot_trajectory_comparison(self, original, reconstructed, title_prefix=""):
+        """Helper method to plot trajectory comparison."""
+        original = original.detach().cpu().numpy()
+        reconstructed = reconstructed.detach().cpu().numpy()
+        x_vals = np.linspace(0, 1.0, original.shape[1])
+        
+        batch_size, time_steps, dof = original.shape
+        
+        for sample_idx in range(batch_size):
+            fig, axes = plt.subplots(dof, 1, figsize=(8, 2 * dof), sharex=True)
+            if dof == 1:
+                axes = [axes]  # Handle single DOF case
+            
+            for i in range(dof):
+                axes[i].plot(x_vals, reconstructed[sample_idx, :, i], 
+                           marker='o', label='Reconstructed', linestyle='-', color='b')
+                axes[i].plot(x_vals, original[sample_idx, :, i], 
+                           marker='*', label='Ground Truth', linestyle='--', color='r')
+                axes[i].set_ylabel(f"DOF {i + 1}")
+                axes[i].grid(True)
+                axes[i].legend(loc="best")
+            
+            axes[-1].set_xlabel("Time (s)")
+            plt.suptitle(f"{title_prefix}Trajectory Comparison - Sample {sample_idx}")
+            plt.tight_layout(rect=[0, 0, 1, 0.96])
+            plt.show()
+    
+    def visualize_reconstruction_error_discrete(self, raw_traj):
+        """Visualize reconstruction error for discrete encoding."""
+        tokens = self.encode_discrete(raw_traj, update_bounds=True)
+        reconstructed = self.decode_discrete(tokens)
+        self._plot_trajectory_comparison(raw_traj, reconstructed, "Discrete ")
+    
+    def visualize_reconstruction_error_continuous(self, raw_traj):
+        """Visualize reconstruction error for continuous encoding."""
+        raw_traj = raw_traj.to(torch.float32)
+        if len(raw_traj.shape) == 2:
+            raw_traj = raw_traj.unsqueeze(0)
+        
+        continuous_tokens = self.encode_continuous(raw_traj, update_bounds=True)
+        reconstructed = self.decode_continuous(continuous_tokens)
+        self._plot_trajectory_comparison(raw_traj, reconstructed, "Continuous ")

bspline_factory.py

@@ -0,0 +1,25 @@
+import torch
+
+from .basis_gn import UniBSplineBasis
+from .uni_bspline import UniformBSpline
+
+
+class SplineFactory:
+
+    @staticmethod
+    def init_splines(mp_type: str,
+                mp_args: dict,
+                num_dof: int = 1,
+                tau: float = 1,
+                dtype: torch.dtype = torch.float32,
+                device: torch.device = "cpu"):
+
+        if mp_type == "uni_bspline":
+            basis_gn = UniBSplineBasis(dtype=dtype, device=device, tau=tau,
+                                       **mp_args)
+            mp = UniformBSpline(basis_gn=basis_gn, num_dof=num_dof,
+                                dtype=dtype, device=device, **mp_args)
+        else:
+            raise NotImplementedError
+
+        return mp

uni_bspline.py

@@ -0,0 +1,462 @@
+
+from typing import Union, Optional
+import logging
+
+import numpy as np
+import torch
+
+from .basis_gn import UniBSplineBasis
+
+class UniformBSpline(torch.nn.Module):
+
+    def __init__(self,
+                 basis_gn: UniBSplineBasis,
+                 num_dof: int,
+                 weights_scale: float = 1.,
+                 dtype: torch.dtype = torch.float32,
+                 device: torch.device = 'cpu',
+                 **kwargs,
+                 ):
+        super().__init__()
+
+        # self.dtype = dtype
+        # self.device = device
+        # batch dim
+        self.add_dim = list()
+
+        self.basis_gn = basis_gn
+        self.num_dof = num_dof
+
+        # Scaling of weights
+        weights_scale = \
+            torch.tensor(weights_scale, dtype=self.dtype, device=self.device)
+        assert weights_scale.ndim <= 1, \
+            "weights_scale should be float or 1-dim vector"
+        self.register_buffer("weights_scale", weights_scale, persistent=False)
+
+        # Value caches
+        # Compute values at these time points
+        self.times = None
+
+        # Learnable parameters
+        self.params = None
+
+        # Initial conditions
+        self.init_pos = None
+        self.init_vel = None
+
+        # Runtime computation results, shall be reset every time when
+        # inputs are reset
+        self.pos = None
+        self.vel = None
+
+
+        #parameters bound
+        # params_bound = kwargs.get("params_bound", None)
+        # if not params_bound:
+        #     params_bound = torch.zeros([2, self.num_params],
+        #                                     dtype=self.dtype,
+        #                                     device=self.device)
+        #     params_bound[0, :] = -torch.inf
+        #     params_bound[1, :] = torch.inf
+        # else:
+        #     params_bound = torch.as_tensor(self.params_bound,
+        #                                         dtype=self.dtype,
+        #                                         device=self.device)
+        # assert list(params_bound.shape) == [2, self.num_params]
+        # self.register_buffer("params_bound", params_bound, persistent=False)
+
+
+        self.end_pos = None
+        self.end_vel = None
+
+        self.params_init = None
+        self.params_end = None
+
+    @property
+    def device(self):
+        return self.basis_gn.device
+
+    @property
+    def dtype(self):
+        return self.basis_gn.dtype
+
+    @property
+    def tau(self):
+        return self.basis_gn.tau
+
+    @property
+    def num_basis(self):
+        return self.basis_gn.num_basis
+
+    @property
+    def num_params(self):
+        return self.basis_gn.num_basis * self.num_dof
+
+    def clear_computation_result(self):
+        """
+        Clear runtime computation result
+
+        Returns:
+            None
+        """
+
+        self.pos = None
+        self.vel = None
+        # also reset tau?
+
+    def set_add_dim(self, add_dim: Union[list, torch.Size]):
+        """
+        Set additional batch dimension
+        Args:
+            add_dim: additional batch dimension
+
+        Returns: None
+
+        """
+        self.add_dim = add_dim
+        self.clear_computation_result()
+
+    def set_times(self, times: Union[torch.Tensor, np.ndarray]):
+        """
+        Set time points
+        Args:
+            times: time points
+
+        Returns:
+            None
+        """
+
+        # Shape of times
+        # [*add_dim, num_times]
+
+        self.times = torch.as_tensor(times, dtype=self.dtype,
+                                     device=self.device)
+        tau = times.reshape(-1)[-1]
+        self.basis_gn.tau.copy_(tau)
+        self.clear_computation_result()
+
+    def set_duration(self, duration: Optional[float], dt: float,):
+        """
+
+        Args:
+            duration: desired duration of trajectory
+            dt: control frequency
+        Returns:
+            None
+        """
+
+        # Shape of times
+        # [*add_dim, num_times]
+        dt = torch.as_tensor(dt, dtype=self.dtype, device=self.device)
+        times = torch.linspace(0, duration, round(duration / dt) + 1,
+                               dtype=self.dtype, device=self.device)
+        times = add_expand_dim(times, list(range(len(self.add_dim))),
+                                    self.add_dim)
+        self.set_times(times)
+
+    def set_params(self,
+                   params: Union[torch.Tensor, np.ndarray]) -> torch.Tensor:
+        """
+        Set MP params
+        Args:
+            params: parameters
+
+        Returns: unused parameters
+
+        """
+        # Shape of params
+        # [*add_dim, num_params]
+
+        params = torch.as_tensor(params, dtype=self.dtype, device=self.device)
+
+        # Check number of params
+        assert params.shape[-1] == self.num_params
+
+        # Set additional batch size
+        self.set_add_dim(list(params.shape[:-1]))
+
+        self.params = params[..., :self.num_params]
+        self.clear_computation_result()
+        return params[..., self.num_params:]
+
+    def update_inputs(self, times=None, params=None,
+                      init_pos=None, init_vel=None, **kwargs):
+
+        if params is not None:
+            self.set_params(params)
+        if times is not None:
+            self.set_times(times)
+        if init_pos is not None:
+            self.set_initial_conditions(init_pos, init_vel, **kwargs)
+
+        end_pos = kwargs.get('end_pos', None)
+        end_vel = kwargs.get('end_vel', None)
+        if any([cond is not None for cond in [end_pos, end_vel]]):
+            self.set_end_condtions(end_pos, end_vel)
+
+    def set_initial_conditions(self,
+                               init_pos: Union[torch.Tensor, np.ndarray],
+                               init_vel: Union[torch.Tensor, np.ndarray],
+                               **kwargs):
+
+        self.init_pos = torch.as_tensor(init_pos, dtype=self.dtype,
+                                        device=self.device)
+        self.init_vel = torch.as_tensor(init_vel, dtype=self.dtype,
+                                   device=self.device) if init_vel is not None else None
+        self.clear_computation_result()
+
+        self.params_init = self.basis_gn.compute_init_params(self.init_pos, self.init_vel)
+        if self.params_init is not None:
+            self.params_init /= self.weights_scale
+
+    def set_end_condtions(self, end_pos: Union[torch.Tensor, np.ndarray],
+                          end_vel: Union[torch.Tensor, np.ndarray], **kwargs):
+        self.end_pos = \
+            torch.as_tensor(end_pos, device=self.device, dtype=self.dtype) \
+                if end_pos is not None else None
+        self.end_vel = \
+            torch.as_tensor(end_vel, device=self.device, dtype=self.dtype) \
+                if end_vel is not None else None
+
+        self.params_end = self.basis_gn.compute_end_params(self.end_pos, self.end_vel)
+        if self.params_end is not None:
+            self.params_end /= self.weights_scale
+
+    def get_traj_pos(self, times=None, params=None,
+                     init_pos=None, init_vel=None, flat_shape=False, **kwargs):
+
+        self.update_inputs(times, params, init_pos, init_vel, **kwargs)
+
+        if self.pos is not None:
+            pos = self.pos
+        else:
+            assert self.params is not None
+
+            # Reshape params
+            # [*add_dim, num_dof * num_basis] -> [*add_dim, num_dof, num_basis]
+            params = self.params.reshape(*self.add_dim, self.num_dof, -1)
+            # extend params with possible init and end conditions
+            # shape: [*add_dim, num_dof, num_ctrlp]
+            if self.params_init is not None:
+                params = torch.cat((self.params_init, params), dim=-1)
+            if self.params_end is not None:
+                params = torch.cat((params, self.params_end), dim=-1)
+
+            # Get basis
+            # Shape: [*add_dim, num_times, num_ctrlp]
+            basis_single_dof = \
+                self.basis_gn.basis(self.times) * self.weights_scale
+
+            # Einsum shape: [*add_dim, num_times, num_ctrlp],
+            #               [*add_dim, num_dof, num_ctrlp]
+            #            -> [*add_dim, num_times, num_dof]
+            pos = torch.einsum('...ik,...jk->...ij', basis_single_dof, params)
+
+            self.pos = pos
+
+        if flat_shape:
+            # Switch axes to [*add_dim, num_dof, num_times]
+            pos = torch.einsum('...ji->...ij', pos)
+
+            # Reshape to [*add_dim, num_dof * num_times]
+            pos = pos.reshape(*self.add_dim, -1)
+
+        return pos
+
+    def get_traj_vel(self, times=None, params=None,
+                     init_pos=None, init_vel=None, flat_shape=False, **kwargs):
+
+        self.update_inputs(times, params, init_pos, init_vel,
+                           **kwargs)
+
+        if self.vel is not None:
+            vel = self.vel
+        else:
+            assert self.params is not None
+
+            # Reshape params
+            # [*add_dim, num_dof * num_basis] -> [*add_dim, num_dof, num_basis]
+            params = self.params.reshape(*self.add_dim, self.num_dof, -1)
+            # extend params with possible init and end conditions
+            # shape: [*add_dim, num_dof, num_ctrlp]
+            if self.params_init is not None:
+                params = torch.cat((self.params_init, params), dim=-1)
+            if self.params_end is not None:
+                params = torch.cat((params, self.params_end), dim=-1)
+
+            # velocity control points shape: [*add_dim, num_dof, num_ctrlp-1]
+            vel_ctrlp = self.basis_gn.velocity_control_points(params)
+            vel_ctrlp = torch.einsum("...ij,...->...ij", vel_ctrlp,
+                                     1 / self.tau)
+
+            # vel_basis shape: [*add_dim, num_times, num_ctrlp-1]
+            vel_basis = self.basis_gn.vel_basis(self.times) * self.weights_scale
+
+            # Einsum shape: [*add_dim, num_times, num_ctrlp-1],
+            #               [*add_dim, num_dof, num_ctrlp-1]
+            #            -> [*add_dim, num_times, num_dof]
+            vel = torch.einsum('...ik,...jk->...ij', vel_basis, vel_ctrlp)
+
+            self.vel = vel
+
+        if flat_shape:
+            # Switch axes to [*add_dim, num_dof, num_times]
+            vel = torch.einsum('...ji->...ij', vel)
+
+            # Reshape to [*add_dim, num_dof * num_times]
+            vel = vel.reshape(*self.add_dim, -1)
+
+        return vel
+
+    def learn_mp_params_from_trajs(self, times: torch.Tensor,
+                                   trajs: torch.Tensor, reg=1e-5, **kwargs):
+
+        # only works for learn_tau=False, learn_delay=False. And delay=0 (or you
+        # need to give the initial condition by yourself)
+
+        # Shape of times
+        # [*add_dim, num_times]
+        #
+        # Shape of trajs:
+        # [*add_dim, num_times, num_dof]
+        #
+        # Shape of params:
+        # [*add_dim, num_dof * num_basis]
+
+        assert trajs.shape[:-1] == times.shape
+        assert trajs.shape[-1] == self.num_dof
+
+        times = torch.as_tensor(times, dtype=self.dtype, device=self.device)
+        trajs = torch.as_tensor(trajs, dtype=self.dtype, device=self.device)
+
+        # Setup stuff
+        self.set_add_dim(list(trajs.shape[:-2]))
+        self.set_times(times)
+        dummy_params = torch.zeros(*self.add_dim, self.num_dof, self.num_basis,
+                                   device=self.device, dtype=self.dtype)
+
+        # Get initial conditions
+        if self.basis_gn.init_cond_order != 0:
+            if any([key in kwargs.keys()
+                    for key in [ "init_pos", "init_vel"]]):
+                logging.warning("uses the given initial conditions")
+                init_pos = kwargs.get("init_pos")
+                init_vel = kwargs.get("init_vel")
+            else:
+                init_pos = trajs[..., 0, :]
+                dt = (times[..., 1] - times[..., 0])
+                init_vel = torch.einsum("...i,...->...i",
+                                        torch.diff(trajs, dim=-2)[..., 0, :],
+                                        1/dt)
+            self.set_initial_conditions(init_pos, init_vel)
+            if self.params_init is not None:
+                dummy_params = torch.cat([self.params_init, dummy_params],
+                                         dim=-1)
+
+        if self.basis_gn.end_cond_order != 0:
+            if any([key in kwargs.keys()
+                    for key in ["end_pos", "end_vel"]]):
+                logging.warning("uses the given end conditions")
+                end_pos = kwargs.get("end_pos")
+                end_vel = kwargs.get("end_vel")
+            else:
+                end_pos = trajs[..., -1, :]
+                dt = (times[..., 1] - times[..., 0])
+                end_vel = torch.einsum("...i,...->...i",
+                                       torch.diff(trajs, dim=-2)[..., -1, :],
+                                       1/dt)
+            self.set_end_condtions(end_pos, end_vel)
+            if self.params_end is not None:
+                dummy_params = torch.cat([dummy_params, self.params_end],
+                                         dim=-1)
+
+        basis_single_dof = self.basis_gn.basis(times) * self.weights_scale
+        # shape: [*add_dim, num_time, num_ctrlp]
+        #        [*add_dim, num_dof, num_ctrlp]
+        #        [*add_dim, num_times, num_dof]
+        pos_det = torch.einsum('...ik,...jk->...ij', basis_single_dof, dummy_params)
+        # swtich axes to [*add_dim, num_dof, num_times]
+        pos_det = torch.einsum('...ij->...ji', pos_det)
+        pos_det = pos_det.reshape(*self.add_dim, -1)
+
+        basis_multi_dofs = self.basis_gn.basis_multi_dofs(self.times, self.num_dof) * self.weights_scale
+        # Solve this: Aw = B -> w = A^{-1} B
+        # Einsum_shape: [*add_dim, num_dof * num_times, num_dof * num_basis]
+        #               [*add_dim, num_dof * num_times, num_dof * num_basis]
+        #            -> [*add_dim, num_dof * num_basis, num_dof * num_basis]
+        A = torch.einsum('...ki,...kj->...ij', basis_multi_dofs,
+                         basis_multi_dofs)
+        A += torch.eye(self.num_params,
+                       dtype=self.dtype,
+                       device=self.device) * reg
+
+        # Swap axis and reshape: [*add_dim, num_times, num_dof]
+        #                     -> [*add_dim, num_dof, num_times]
+        trajs = torch.einsum("...ij->...ji", trajs)
+        # Reshape [*add_dim, num_dof, num_times]
+        #      -> [*add_dim, num_dof * num_times]
+        trajs = trajs.reshape([*self.add_dim, -1])
+
+        # Position minus initial condition terms,
+        pos_w = trajs - pos_det
+
+        # Einsum_shape: [*add_dim, num_dof * num_times, num_dof * num_basis]
+        #               [*add_dim, num_dof * num_times]
+        #            -> [*add_dim, num_dof * num_basis]
+        B = torch.einsum('...ki,...k->...i', basis_multi_dofs, pos_w)
+
+        # Shape of weights: [*add_dim, num_dof * num_basis]
+        params = torch.linalg.solve(A, B)
+
+        self.set_params(params)
+
+        return {"params": params,
+                "init_pos": self.init_pos,
+                "init_vel": self.init_vel,
+                "end_pos": self.end_pos,
+                "end_vel": self.end_vel,
+                }
+
+
+def add_expand_dim(data: Union[torch.Tensor, np.ndarray],
+                   add_dim_indices: [int],
+                   add_dim_sizes: [int]) -> Union[torch.Tensor, np.ndarray]:
+    """
+    Add additional dimensions to tensor and expand accordingly
+    Args:
+        data: tensor to be operated. Torch.Tensor or numpy.ndarray
+        add_dim_indices: the indices of added dimensions in the result tensor
+        add_dim_sizes: the expanding size of the additional dimensions
+
+    Returns:
+        result: result tensor after adding and expanding
+    """
+    num_data_dim = data.ndim
+    num_dim_to_add = len(add_dim_indices)
+
+    add_dim_reverse_indices = [num_data_dim + num_dim_to_add + idx for idx in
+                               add_dim_indices]
+
+    str_add_dim = ""
+    str_expand = ""
+    add_dim_index = 0
+    for dim in range(num_data_dim + num_dim_to_add):
+        if dim in add_dim_indices or dim in add_dim_reverse_indices:
+            str_add_dim += "None, "
+            str_expand += str(add_dim_sizes[add_dim_index]) + ", "
+            add_dim_index += 1
+        else:
+            str_add_dim += ":, "
+            if type(data) == torch.Tensor:
+                str_expand += "-1, "
+            elif type(data) == np.ndarray:
+                str_expand += "1, "
+            else:
+                raise NotImplementedError
+
+    str_add_dime_eval = "data[" + str_add_dim + "]"
+    if type(data) == torch.Tensor:
+        return eval("eval(str_add_dime_eval).expand(" + str_expand + ")")
+    else:
+        return eval("np.tile(eval(str_add_dime_eval),[" + str_expand + "])")

utils.py

@@ -0,0 +1,166 @@
+import torch
+import einops
+from typing import Union
+
+def continuous_to_discrete(tensor, min_val=None, max_val=None, num_bins=256):
+    """
+    Convert a continuous PyTorch tensor to discrete tokens in the range [0, 255].
+
+    Args:
+        tensor (torch.Tensor): Input tensor with continuous values.
+        min_val (float, optional): Minimum value for normalization. If None, use tensor.min().
+        max_val (float, optional): Maximum value for normalization. If None, use tensor.max().
+
+    Returns:
+        torch.Tensor: Discretized tensor with values in the range [0, 255].
+    """
+
+    if min_val is None:
+        min_val = tensor.min()
+    if max_val is None:
+        max_val = tensor.max()
+
+    # Normalize the tensor to [0, 1]
+    assert torch.all(tensor >= min_val - 1e-3), "Input tensor has values below min_val"
+    assert torch.all(tensor <= max_val + 1e-3), "Input tensor has values above max_val"
+    normalized_tensor = (tensor - min_val) / (max_val - min_val)
+    normalized_tensor = torch.clamp(normalized_tensor, 0, 1)
+
+    # Ensure no out-of-bound values
+    # Scale to [0, 255] and quantize to integers
+    discrete_tensor = torch.round(normalized_tensor * (num_bins-1)).to(torch.long)
+    return discrete_tensor
+
+
+def discrete_to_continuous(discrete_tensor, min_val=0, max_val=1, num_bins=256):
+    """
+    Convert a discrete PyTorch tensor with values in the range [0, 255]
+    back to continuous values in the range [min_val, max_val].
+
+    Args:
+        discrete_tensor (torch.Tensor): Input tensor with discrete values (0 to 255).
+        min_val (float): Minimum value of the original continuous range.
+        max_val (float): Maximum value of the original continuous range.
+
+    Returns:
+        torch.Tensor: Continuous tensor with values in the range [min_val, max_val].
+    """
+    # Map discrete tokens to [0, 1]
+    # Normalize the tensor to [0, 1]
+    normalized_tensor = discrete_tensor.float() / (num_bins-1)
+
+    # Map normalized values to [min_val, max_val]
+    continuous_tensor = normalized_tensor * (max_val - min_val) + min_val
+
+    # Ensure no out-of-bound values
+    continuous_tensor = torch.clamp(continuous_tensor, min_val, max_val)
+    return continuous_tensor
+
+
+def normalize_tensor(tensor, w_min, w_max, norm_min=-1.0, norm_max=1.0):
+    """
+    Normalize a tensor from its original range [w_min, w_max] to a new range [norm_min, norm_max].
+    
+    Args:
+        tensor (torch.Tensor): Input tensor to be normalized
+        w_min (float): Minimum value bound of the original tensor
+        w_max (float): Maximum value bound of the original tensor
+        norm_min (float, optional): Minimum value of the normalized range. Defaults to 0.0.
+        norm_max (float, optional): Maximum value of the normalized range. Defaults to 1.0.
+    
+    Returns:
+        torch.Tensor: Normalized tensor with values in range [norm_min, norm_max]
+    """
+    
+    # Clip the input tensor to be within [w_min, w_max]
+    clipped_tensor = torch.clamp(tensor, w_min, w_max)
+    
+    # Normalize to [0, 1] range first
+    normalized = (clipped_tensor - w_min) / (w_max - w_min)
+    
+    # Scale to the desired [norm_min, norm_max] range
+    normalized = normalized * (norm_max - norm_min) + norm_min
+    
+    return normalized
+
+def denormalize_tensor(normalized_tensor, w_min, w_max, norm_min=-1.0, norm_max=1.0):
+    """
+    Denormalize a tensor from the normalized range [norm_min, norm_max] back to the original range [w_min, w_max].
+    
+    Args:
+        normalized_tensor (torch.Tensor): Normalized input tensor
+        w_min (float): Minimum value bound of the original range
+        w_max (float): Maximum value bound of the original range
+        norm_min (float, optional): Minimum value of the normalized range. Defaults to 0.0.
+        norm_max (float, optional): Maximum value of the normalized range. Defaults to 1.0.
+    
+    Returns:
+        torch.Tensor: Denormalized tensor with values in range [w_min, w_max]
+    """
+    
+    # Clip the normalized tensor to be within [norm_min, norm_max]
+    clipped_tensor = torch.clamp(normalized_tensor, norm_min, norm_max)
+    
+    # Scale from [norm_min, norm_max] to [0, 1] first
+    denormalized = (clipped_tensor - norm_min) / (norm_max - norm_min)
+    
+    # Scale to the original [w_min, w_max] range
+    denormalized = denormalized * (w_max - w_min) + w_min
+    
+    return denormalized
+
+
+def tensor_linspace(start: Union[float, int, torch.Tensor],
+                    end: Union[float, int, torch.Tensor],
+                    steps: int) -> torch.Tensor:
+    """
+    Vectorized version of torch.linspace.
+    Modified from:
+    https://github.com/zhaobozb/layout2im/blob/master/models/bilinear.py#L246
+
+    Args:
+        start: start value, scalar or tensor
+        end: end value, scalar or tensor
+        steps: num of steps
+
+    Returns:
+        linspace tensor
+    """
+    # Shape of start:
+    # [*add_dim, dim_data] or a scalar
+    #
+    # Shape of end:
+    # [*add_dim, dim_data] or a scalar
+    #
+    # Shape of out:
+    # [*add_dim, steps, dim_data]
+
+    # - out: Tensor of shape start.size() + (steps,), such that
+    #   out.select(-1, 0) == start, out.select(-1, -1) == end,
+    #   and the other elements of out linearly interpolate between
+    #   start and end.
+
+    if isinstance(start, torch.Tensor) and not isinstance(end, torch.Tensor):
+        end += torch.zeros_like(start)
+    elif not isinstance(start, torch.Tensor) and isinstance(end, torch.Tensor):
+        start += torch.zeros_like(end)
+    elif isinstance(start, torch.Tensor) and isinstance(end, torch.Tensor):
+        assert start.size() == end.size()
+    else:
+        return torch.linspace(start, end, steps)
+
+    view_size = start.size() + (1,)
+    w_size = (1,) * start.dim() + (steps,)
+    out_size = start.size() + (steps,)
+
+    start_w = torch.linspace(1, 0, steps=steps).to(start)
+    start_w = start_w.view(w_size).expand(out_size)
+    end_w = torch.linspace(0, 1, steps=steps).to(start)
+    end_w = end_w.view(w_size).expand(out_size)
+
+    start = start.contiguous().view(view_size).expand(out_size)
+    end = end.contiguous().view(view_size).expand(out_size)
+
+    out = start_w * start + end_w * end
+    out = torch.einsum('...ji->...ij', out)
+    return out