Fix seed shape issue

src/curobo/cuda_robot_model/cuda_robot_generator.py

@@ -969,11 +969,7 @@ def _update_joint_limits(self):
                 self.cspace.max_acceleration.unsqueeze(0),
             ]
         )
-        # clip joint position:
         # NOTE: change this to be per joint
-        joint_limits["position"][0] += self.cspace.position_limit_clip
-        joint_limits["position"][1] -= self.cspace.position_limit_clip
-        joint_limits["velocity"][0] *= self.cspace.velocity_scale
-        joint_limits["velocity"][1] *= self.cspace.velocity_scale
-
         self._joint_limits = JointLimits(joint_names=self.joint_names, **joint_limits)
+        self._joint_limits.clip_position_Limits(self.cspace.position_limit_clip)
+        self._joint_limits.apply_scale(velocity_scale=self.cspace.velocity_scale)

src/curobo/cuda_robot_model/types.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 # Standard Library
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
@@ -53,6 +53,11 @@ class JointLimits:
     effort: Optional[torch.Tensor] = None
     tensor_args: TensorDeviceType = TensorDeviceType()
 
+    position_clipped: bool = field(default=False, init=False)
+    velocity_scaled: bool = field(default=False, init=False)
+    acceleration_scaled: bool = field(default=False, init=False)
+    jerk_scaled: bool = field(default=False, init=False)
+
     @staticmethod
     def from_data_dict(data: Dict, tensor_args: TensorDeviceType = TensorDeviceType()):
         p = tensor_args.to_device(data["position"])
@@ -84,6 +89,23 @@ def copy_(self, new_jl: JointLimits):
         self.effort = copy_if_not_none(new_jl.effort, self.effort)
         return self
 
+    def clip_position_Limits(self, position_clip: float):
+        if not self.position_clipped:
+            # Clip the position limits
+            self.position[0] += position_clip
+            self.position[1] -= position_clip
+
+    def apply_scale(self, velocity_scale: Optional[float] = None, acceleration_scale: Optional[float] = None, jerk_scale: Optional[float] = None):
+        if velocity_scale is not None and not self.velocity_scaled:
+            self.velocity = self.velocity * velocity_scale
+            self.velocity_scaled = True
+        if acceleration_scale is not None and not self.acceleration_scaled:
+            self.acceleration = self.acceleration * acceleration_scale
+            self.acceleration_scaled = True
+        if jerk_scale is not None and not self.jerk_scaled:
+            self.jerk = self.jerk * jerk_scale
+            self.jerk_scaled = True
+
 
 @dataclass
 class CSpaceConfig:
@@ -196,16 +218,6 @@ def clone(self) -> CSpaceConfig:
             jerk_scale=self.jerk_scale.clone(),
         )
 
-    def scale_joint_limits(self, joint_limits: JointLimits):
-        if self.velocity_scale is not None:
-            joint_limits.velocity = joint_limits.velocity * self.velocity_scale
-        if self.acceleration_scale is not None:
-            joint_limits.acceleration = joint_limits.acceleration * self.acceleration_scale
-        if self.jerk_scale is not None:
-            joint_limits.jerk = joint_limits.jerk * self.jerk_scale
-
-        return joint_limits
-
     @staticmethod
     def load_from_joint_limits(
         joint_position_upper: torch.Tensor,
@@ -315,7 +327,7 @@ def __post_init__(self):
         if self.cspace is None and self.joint_limits is not None:
             self.load_cspace_cfg_from_kinematics()
         if self.joint_limits is not None and self.cspace is not None:
-            self.joint_limits = self.cspace.scale_joint_limits(self.joint_limits)
+            self.joint_limits.apply_scale(self.cspace.velocity_scale, self.cspace.acceleration_scale, self.cspace.jerk_scale)
         if self.link_spheres is not None and self.reference_link_spheres is None:
             self.reference_link_spheres = self.link_spheres.clone()
 

src/curobo/rollout/cost/bound_cost.py

@@ -37,7 +37,7 @@ class BoundCostType(Enum):
     BOUNDS_SMOOTH = 2
 
 
-@dataclass
+@dataclass(eq=False)
 class BoundCostConfig(CostConfig):
     joint_limits: Optional[JointLimits] = None
     smooth_weight: Optional[List[float]] = None
@@ -298,12 +298,15 @@ def forward_all_bound_cost(
     p,
     v,
     a,
+    j,
     p_lower_bounds,
     p_upper_bounds,
     v_lower_bounds,
     v_upper_bounds,
     a_lower_bounds,
     a_upper_bounds,
+    j_lower_bounds,
+    j_upper_bounds,
     weight,
 ):
     # c = torch.sum(
@@ -325,7 +328,9 @@ def forward_all_bound_cost(
             + torch.nn.functional.relu(v - v_upper_bounds)
             + torch.nn.functional.relu(a_lower_bounds - a)
             + torch.nn.functional.relu(a - a_upper_bounds)
-        ),
+            + torch.nn.functional.relu(j_lower_bounds - j)
+            + torch.nn.functional.relu(j - j_upper_bounds)
+            ),
         dim=-1,
     )
 
@@ -661,7 +666,7 @@ def forward_bound_pos_warp(
 
     w = wp.float32(0.0)
     c_p = wp.float32(0.0)
-    g_p = wp.float32(0.0)
+    g_p = wp.float32(0.0) # gradient for position
     c_total = wp.float32(0.0)
 
     # we launch batch * horizon * dof kernels
@@ -670,12 +675,13 @@ def forward_bound_pos_warp(
     d_id = tid - (b_id * horizon * dof + h_id * dof)
     if b_id >= batch_size or h_id >= horizon or d_id >= dof:
         return
-
-    # read weights:
+    # read activation distance (a ratio):
     eta_p = activation_distance[0]
+    # read weights:
     w = weight[0]
 
     n_w = wp.float32(0.0)
+    # read null space weight:
     n_w = null_weight[0]
     target_p = wp.float32(0.0)
     target_id = wp.int32(0.0)
@@ -684,11 +690,14 @@ def forward_bound_pos_warp(
         target_id = retract_idx[b_id]
         target_id = target_id * dof + d_id
         target_p = retract_config[target_id]
+    # read lower and upper bounds:
     p_l = p_b[d_id]
     p_u = p_b[dof + d_id]
 
+    # compute activation distance as ratio of limits range:
     p_range = p_u - p_l
     eta_p = eta_p * (p_range)
+    # shrink lower and upper limits by activation distance:
     p_l += eta_p
     p_u -= eta_p
     if p_range < 1.0:
@@ -697,17 +706,16 @@ def forward_bound_pos_warp(
     # compute cost:
     b_addrs = b_id * horizon * dof + h_id * dof + d_id
 
-    # read buffers:
-
+    # read current position:
     c_p = pos[b_addrs]
 
+    # compute null space cost and its gradient:
     if n_w > 0.0:
         error = c_p - target_p
         c_total = n_w * error * error
         g_p = 2.0 * n_w * error
 
     # bound cost:
-
     delta = 0.0
     if c_p < p_l:
         delta = c_p - p_l
@@ -756,18 +764,25 @@ def forward_bound_warp(
     d_id = wp.int32(0)
     b_addrs = int(0)
 
-    w = wp.float32(0.0)
+    c_p = wp.float32(0.0)
     c_v = wp.float32(0.0)
     c_a = wp.float32(0.0)
-    c_p = wp.float32(0.0)
+    c_j = wp.float32(0.0)
+
     g_p = wp.float32(0.0)
     g_v = wp.float32(0.0)
     g_a = wp.float32(0.0)
-    b_wv = float(0.0)
-    b_wa = float(0.0)
-    b_wj = float(0.0)
+    g_j = wp.float32(0.0)
+
+    b_wp = wp.float32(0.0)
+    b_wv = wp.float32(0.0)
+    b_wa = wp.float32(0.0)
+    b_wj = wp.float32(0.0)
+
     c_total = wp.float32(0.0)
 
+    delta = float(0.0)
+
     # we launch batch * horizon * dof kernels
     b_id = tid / (horizon * dof)
     h_id = (tid - (b_id * horizon * dof)) / dof
@@ -785,7 +800,7 @@ def forward_bound_warp(
         target_p = retract_config[target_id]
 
     # read weights:
-    w = weight[0]
+    b_wp = weight[0]
     b_wv = weight[1]
     b_wa = weight[2]
     b_wj = weight[3]
@@ -794,18 +809,17 @@ def forward_bound_warp(
     b_addrs = b_id * horizon * dof + h_id * dof + d_id
 
     # read buffers:
+    c_p = pos[b_addrs]
     c_v = vel[b_addrs]
     c_a = acc[b_addrs]
-    c_p = pos[b_addrs]
+    c_j = jerk[b_addrs]
 
     # if w_j > 0.0:
     eta_p = activation_distance[0]
     eta_v = activation_distance[1]
     eta_a = activation_distance[2]
     eta_j = activation_distance[3]
 
-    c_j = jerk[b_addrs]
-
     p_l = p_b[d_id]
     p_u = p_b[dof + d_id]
     p_range = p_u - p_l
@@ -834,7 +848,17 @@ def forward_bound_warp(
     j_l += eta_j
     j_u -= eta_j
 
-    delta = float(0.0)
+    # TODO: Is this necessary?
+    # TODO: Why is sometimes (1/range) used while others (1 / range**2)
+    if p_range < 1.0:
+        b_wp = (1.0 / p_range) * b_wp
+    if v_range < 1.0:
+        b_wv = b_wv * (1.0 / v_range)
+    if a_range < 1.0:
+        b_wa = b_wa * (1.0 / a_range)
+    if j_range < 1.0:
+        b_wj = b_wj * (1.0 / j_range)
+
     if n_w > 0.0:
         error = c_p - target_p
         c_total = n_w * error * error
@@ -848,8 +872,8 @@ def forward_bound_warp(
     elif c_p > p_u:
         delta = c_p - p_u
 
-    c_total += w * delta * delta
-    g_p += 2.0 * w * delta
+    c_total += b_wp * delta * delta
+    g_p += 2.0 * b_wp * delta
 
     # bound cost:
     delta = 0.0

src/curobo/rollout/cost/cost_base.py

@@ -19,7 +19,7 @@
 from curobo.types.base import TensorDeviceType
 
 
-@dataclass
+@dataclass(eq=False)
 class CostConfig:
     weight: Union[torch.Tensor, float, List[float]]
     tensor_args: TensorDeviceType = None

src/curobo/rollout/cost/primitive_collision_cost.py

@@ -23,7 +23,7 @@
 from curobo.util.torch_utils import get_torch_jit_decorator
 
 
-@dataclass
+@dataclass(eq=False)
 class PrimitiveCollisionCostConfig(CostConfig):
     """Create Collision Cost Configuration."""
 

src/curobo/rollout/cost/self_collision_cost.py

@@ -23,7 +23,7 @@
 from .cost_base import CostBase, CostConfig
 
 
-@dataclass
+@dataclass(eq=False)
 class SelfCollisionCostConfig(CostConfig):
     self_collision_kin_config: Optional[SelfCollisionKinematicsConfig] = None
 

src/curobo/types/state.py

@@ -193,6 +193,13 @@ def get_state_tensor(self):
         state_tensor = torch.cat((self.position, velocity, acceleration, jerk), dim=-1)
         return state_tensor
 
+    @property
+    def augmented_state(self):
+        velocity = self.velocity
+        if velocity is None:
+            velocity = self.position * 0.0
+        return torch.cat((self.position, velocity), dim=-1)
+
     @staticmethod
     def from_state_tensor(state_tensor, joint_names=None, dof=7):
         return JointState(
@@ -371,6 +378,18 @@ def zeros(
             joint_names=joint_names,
         )
 
+    def contiguous(self):
+        p = v = a = j = None
+        if self.position is not None:
+            p = self.position.contiguous()
+        if self.velocity is not None:
+            v = self.velocity.contiguous()
+        if self.acceleration is not None:
+            a = self.acceleration.contiguous()
+        if self.jerk is not None:
+            j = self.jerk.contiguous()
+        return JointState(p, v, a, self.joint_names, jerk=j)
+
     def detach(self):
         self.position = self.position.detach()
         self.velocity = self.velocity.detach()