fixed bug with the collision objective. Created fast collision objective. Also added this to the original library.

app.py

@@ -30,6 +30,10 @@ from jax_ik.objectives import (
     DistanceObjTraj,
     SphereCollisionPenaltyObjTraj,
 )
+try:
+    from collision_objectives import FastSphereCollisionPenaltyObjTraj as _FastColl
+except Exception:  # fallback if file missing
+    _FastColl = SphereCollisionPenaltyObjTraj
 
 def download_and_setup_files():
     os.makedirs("files", exist_ok=True)
@@ -78,13 +82,6 @@ class IKServer:
         self.max_num_steps_global = 20000
         self.max_learning_rate = 5.0
 
-        # Caches for solver reuse (LRU)
-        self.solver_cache = {}
-        self.urdf_solver_cache = {}
-        self.cache_access_order = []
-        self.urdf_cache_access_order = []
-        self.max_cache_size = 5
-
         # Animation buffers
         self.animation_frames_agent = []
         self.animation_frames_urdf = []
@@ -145,23 +142,6 @@ class IKServer:
         allowed = self.urdf_available_bones if is_urdf else self.available_bones
         return [b for b in bones if b in allowed]
 
-    def _cache_key(self, bones, num_steps):
-        return tuple(sorted(bones)) + (int(num_steps),)
-
-    def _evict_lru(self, is_urdf=False):
-        cache = self.urdf_solver_cache if is_urdf else self.solver_cache
-        order = self.urdf_cache_access_order if is_urdf else self.cache_access_order
-        if not order:
-            return
-        k = order.pop(0)
-        cache.pop(k, None)
-        gc.collect()
-
-    def _invalidate_caches(self):
-        self.solver_cache.clear(); self.urdf_solver_cache.clear()
-        self.cache_access_order.clear(); self.urdf_cache_access_order.clear()
-        gc.collect()
-
     def _create_solver(self, bones, is_urdf, num_steps):
         if is_urdf:
             return InverseKinematicsSolver(
@@ -188,29 +168,6 @@ class IKServer:
             compute_sdf=False,
         )
 
-    def _get_solver(self, bones, is_urdf=False, num_steps=None):
-        if num_steps is None:
-            num_steps = self.args.num_steps
-        num_steps = self._sanitize_num_steps(num_steps)
-        key = self._cache_key(bones, num_steps)
-        cache = self.urdf_solver_cache if is_urdf else self.solver_cache
-        order = self.urdf_cache_access_order if is_urdf else self.cache_access_order
-        if key in cache:
-            if key in order: order.remove(key)
-            order.append(key)
-            return cache[key]
-        if len(cache) >= self.max_cache_size:
-            self._evict_lru(is_urdf)
-        try:
-            solver = self._create_solver(bones, is_urdf, num_steps)
-        except Exception as e:
-            logger.warning(f"Primary solver creation failure (is_urdf={is_urdf}): {e}; using defaults")
-            bones = self.urdf_default_controlled_bones if is_urdf else self.default_controlled_bones
-            solver = self._create_solver(bones, is_urdf, num_steps)
-        cache[key] = solver
-        order.append(key)
-        return solver
-
     def _rebuild_solver_safe(self, is_urdf=False, force_defaults=False):
         try:
             if is_urdf:
@@ -232,8 +189,7 @@ class IKServer:
             return False
 
     def _attempt_solver_recovery(self, is_urdf=False):
-        logger.warning("Attempting full solver recovery (invalidate caches + defaults)")
-        self._invalidate_caches()
+        logger.warning("Attempting full solver recovery (reset to defaults)")
         if is_urdf:
             self.urdf_current_controlled_bones = self.urdf_default_controlled_bones.copy()
             return self._rebuild_solver_safe(True, force_defaults=True)
@@ -291,7 +247,8 @@ class IKServer:
         self.selectable_bones = [b for b in self.available_bones if b in self.bounds_dict]
         self.current_controlled_bones = self.default_controlled_bones.copy()
         self.current_end_effector = self.default_end_effector
-        self.solver = self._get_solver(self.current_controlled_bones, is_urdf=False, num_steps=self.current_num_steps)
+        # Direct solver creation (no cache)
+        self.solver = self._create_solver(self.current_controlled_bones, False, self.current_num_steps)
         self.initial_rotations = np.zeros(len(self.solver.controlled_bones) * 3, dtype=np.float32)
         self.best_angles = self.initial_rotations.copy()
         self.mesh_data = load_mesh_data_from_gltf(self.args.gltf_file, self.solver.fk_solver)
@@ -314,7 +271,7 @@ class IKServer:
         self.urdf_selectable_bones = list(self.urdf_available_bones)
         self.urdf_current_controlled_bones = self.urdf_default_controlled_bones.copy()
         self.urdf_current_end_effector = self.urdf_default_end_effector
-        self.urdf_solver = self._get_solver(self.urdf_current_controlled_bones, is_urdf=True, num_steps=self.urdf_current_num_steps)
+        self.urdf_solver = self._create_solver(self.urdf_current_controlled_bones, True, self.urdf_current_num_steps)
         self.urdf_initial_rotations = np.zeros(len(self.urdf_solver.controlled_bones) * 3, dtype=np.float32)
         self.urdf_best_angles = self.urdf_initial_rotations.copy()
         self.urdf_mesh_data = load_mesh_data_from_urdf(self.urdf_file, self.urdf_solver.fk_solver)
@@ -328,7 +285,9 @@ class IKServer:
             use_head=True,
             weight=1.0,
         )
-        self.collision_obj = SphereCollisionPenaltyObjTraj(
+        use_fast = os.environ.get("USE_FAST_COLLISION", "1") != "0"
+        CollCls = _FastColl if use_fast else SphereCollisionPenaltyObjTraj
+        self.collision_obj = CollCls(
             {"center":[0.1,0.0,0.35],"radius":0.1},
             min_clearance=0.0,
             weight=1.0,
@@ -345,7 +304,9 @@ class IKServer:
             use_head=True,
             weight=1.0,
         )
-        self.urdf_collision_obj = SphereCollisionPenaltyObjTraj(
+        use_fast = os.environ.get("USE_FAST_COLLISION", "1") != "0"
+        CollCls = _FastColl if use_fast else SphereCollisionPenaltyObjTraj
+        self.urdf_collision_obj = CollCls(
             {"center":[0.2,0.0,0.35],"radius":0.1},
             min_clearance=0.0,
             weight=1.0,
@@ -359,28 +320,23 @@ class IKServer:
         tgt = np.array(payload.get("target",[0.0,0.2,0.35]))
         self.distance_obj.update_params({"bone_name": self.current_end_effector, "target_points": tgt, "weight": float(payload.get("distance_weight",1.0))})
         # collision updates (weight, center, radius, min_clearance)
-        if payload.get("collision_enabled", False):
-            coll_update = {"weight": float(payload.get("collision_weight",1.0))}
-            center = payload.get("collision_center")
-            if isinstance(center, (list, tuple)) and len(center)==3:
-                coll_update["center"] = center
-            radius = payload.get("collision_radius")
-            if radius is not None:
-                coll_update["radius"] = float(radius)
-            if "collision_min_clearance" in payload:
-                coll_update["min_clearance"] = float(payload.get("collision_min_clearance", 0.0))
-            self.collision_obj.update_params(coll_update)
-        else:
-            # still update weight / min_clearance so they can be changed next enable
-            mc = payload.get("collision_min_clearance")
-            upd = {"weight": float(payload.get("collision_weight",1.0))}
-            if mc is not None:
-                upd["min_clearance"] = float(mc)
-            self.collision_obj.update_params(upd)
+        collision_enabled = payload.get("collision_enabled", False)
+        coll_update = {"weight": float(payload.get("collision_weight",1.0)) if collision_enabled else 0.0}
+        center = payload.get("collision_center")
+        if isinstance(center, (list, tuple)) and len(center)==3:
+            coll_update["center"] = center
+        radius = payload.get("collision_radius")
+        if radius is not None:
+            coll_update["radius"] = float(radius)
+        if "collision_min_clearance" in payload:
+            coll_update["min_clearance"] = float(payload.get("collision_min_clearance", 0.0))
+        # Always update params so compiled graph stays stable
+        self.collision_obj.update_params(coll_update)
         subpoints = self._safe_int(payload.get("subpoints",1), 1, 1, 100)
         mandatory, optional = [], []
         if payload.get("distance_enabled", True): mandatory.append(self.distance_obj)
-        if payload.get("collision_enabled", False): optional.append(self.collision_obj)
+        # Always include collision objective to avoid JIT retrace when toggling
+        optional.append(self.collision_obj)
         if payload.get("bone_zero_enabled", True):
             optional.append(BoneZeroRotationObj(weight=float(payload.get("bone_zero_weight",0.05))))
         if payload.get("derivative_enabled", True) and subpoints > 1:
@@ -421,27 +377,22 @@ class IKServer:
     def _build_urdf_objectives(self, payload):
         tgt = np.array(payload.get("target",[0.3,0.3,0.35]))
         self.urdf_distance_obj.update_params({"bone_name": self.urdf_current_end_effector, "target_points": tgt, "weight": float(payload.get("distance_weight",1.0))})
-        if payload.get("collision_enabled", False):
-            coll_update = {"weight": float(payload.get("collision_weight",1.0))}
-            center = payload.get("collision_center")
-            if isinstance(center, (list, tuple)) and len(center)==3:
-                coll_update["center"] = center
-            radius = payload.get("collision_radius")
-            if radius is not None:
-                coll_update["radius"] = float(radius)
-            if "collision_min_clearance" in payload:
-                coll_update["min_clearance"] = float(payload.get("collision_min_clearance", 0.0))
-            self.urdf_collision_obj.update_params(coll_update)
-        else:
-            mc = payload.get("collision_min_clearance")
-            upd = {"weight": float(payload.get("collision_weight",1.0))}
-            if mc is not None:
-                upd["min_clearance"] = float(mc)
-            self.urdf_collision_obj.update_params(upd)
+        collision_enabled = payload.get("collision_enabled", False)
+        coll_update = {"weight": float(payload.get("collision_weight",1.0)) if collision_enabled else 0.0}
+        center = payload.get("collision_center")
+        if isinstance(center, (list, tuple)) and len(center)==3:
+            coll_update["center"] = center
+        radius = payload.get("collision_radius")
+        if radius is not None:
+            coll_update["radius"] = float(radius)
+        if "collision_min_clearance" in payload:
+            coll_update["min_clearance"] = float(payload.get("collision_min_clearance", 0.0))
+        self.urdf_collision_obj.update_params(coll_update)
         subpoints = self._safe_int(payload.get("subpoints",1),1,1,100)
         mandatory, optional = [], []
         if payload.get("distance_enabled", True): mandatory.append(self.urdf_distance_obj)
-        if payload.get("collision_enabled", False): optional.append(self.urdf_collision_obj)
+        # Always include collision objective with weight possibly zero
+        optional.append(self.urdf_collision_obj)
         if payload.get("bone_zero_enabled", True):
             optional.append(BoneZeroRotationObj(weight=float(payload.get("bone_zero_weight",0.05))))
         if payload.get("derivative_enabled", True) and subpoints > 1:
@@ -481,7 +432,7 @@ class IKServer:
                     self.current_controlled_bones = bones
                     self.current_end_effector = eff
                     self.current_num_steps = int(num_steps)
-                    self.solver = self._get_solver(bones, False, self.current_num_steps)
+                    self.solver = self._create_solver(bones, False, self.current_num_steps)
                     self.initial_rotations = np.zeros(len(self.solver.controlled_bones)*3, dtype=np.float32)
                     self.best_angles = self.initial_rotations.copy()
                     self._setup_agent_objectives()
@@ -507,7 +458,7 @@ class IKServer:
                     self.urdf_current_controlled_bones = bones
                     self.urdf_current_end_effector = eff
                     self.urdf_current_num_steps = int(num_steps)
-                    self.urdf_solver = self._get_solver(bones, True, self.urdf_current_num_steps)
+                    self.urdf_solver = self._create_solver(bones, True, self.urdf_current_num_steps)
                     self.urdf_initial_rotations = np.zeros(len(self.urdf_solver.controlled_bones)*3, dtype=np.float32)
                     self.urdf_best_angles = self.urdf_initial_rotations.copy()
                     self._setup_urdf_objectives()
@@ -639,11 +590,7 @@ class IKServer:
             return
         try:
             gc.collect()
-            # Memory guard: if RSS > 2GB, clear caches
-            rss = self.process.memory_info().rss
-            if rss > 16 * 1024**3:
-                logger.warning("High memory usage detected; invalidating caches")
-                self._invalidate_caches()
+            # Removed cache invalidation logic (no caches anymore)
         except Exception:
             pass
         self.last_cleanup_time = now
@@ -757,10 +704,9 @@ class IKServer:
         @self.app.post("/reset")
         def reset():
             with self.solve_lock:
-                self._invalidate_caches()
                 self._attempt_solver_recovery(False)
                 self._attempt_solver_recovery(True)
-                return {"status":"ok","message":"solvers & caches reset"}
+                return {"status":"ok","message":"solvers reset"}
 
         @self.app.get("/health")
         def health():
@@ -769,8 +715,6 @@ class IKServer:
                 "status":"ok",
                 "agent_frames": len(self.animation_frames_agent),
                 "urdf_frames": len(self.animation_frames_urdf),
-                "cache_agent": len(self.solver_cache),
-                "cache_urdf": len(self.urdf_solver_cache),
                 "agent_fail_count": self.agent_fail_count,
                 "urdf_fail_count": self.urdf_fail_count,
                 "last_agent_error": self.last_agent_error,

collision_objectives.py

@@ -0,0 +1,83 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax.tree_util import register_pytree_node_class
+from jax_ik.objectives import ObjectiveFunction
+
+@register_pytree_node_class
+class FastSphereCollisionPenaltyObjTraj(ObjectiveFunction):
+    """Vectorized sphere collision penalty over bone segments.
+
+    Weight stored as Python float (static aux) to avoid tracer-to-Python
+    concretization when objective always present with varying weight.
+    """
+    def __init__(self, sphere_collider: dict, weight: float = 1.0, min_clearance: float = 0.05, segment_radius: float = 0.02):
+        self.center = jnp.asarray(sphere_collider["center"], jnp.float32)
+        self.radius = jnp.asarray(sphere_collider["radius"], jnp.float32)
+        self.min_clearance = jnp.asarray(min_clearance, jnp.float32)
+        self.segment_radius = jnp.asarray(segment_radius, jnp.float32)
+        self.weight = float(weight)
+
+    # pytree impl ------------------------------------------------------------
+    def tree_flatten(self):
+        # weight treated as static (aux) so changing it may retrace but avoids concretization errors
+        return (self.center, self.radius, self.min_clearance, self.segment_radius), (self.weight,)
+
+    @classmethod
+    def tree_unflatten(cls, aux, leaves):
+        (weight,) = aux
+        c, r, mc, sr = leaves
+        return cls(dict(center=c, radius=r), weight, mc, sr)
+
+    # API --------------------------------------------------------------------
+    def update_params(self, p: dict) -> None:
+        if "sphere_collider" in p:
+            collider = p["sphere_collider"]
+            if "center" in collider:
+                self.center = jnp.asarray(collider["center"], jnp.float32)
+            if "radius" in collider:
+                self.radius = jnp.asarray(collider["radius"], jnp.float32)
+        if "center" in p:
+            self.center = jnp.asarray(p["center"], jnp.float32)
+        if "radius" in p:
+            self.radius = jnp.asarray(p["radius"], jnp.float32)
+        if "min_clearance" in p:
+            self.min_clearance = jnp.asarray(p["min_clearance"], jnp.float32)
+        if "segment_radius" in p:
+            self.segment_radius = jnp.asarray(p["segment_radius"], jnp.float32)
+        if "weight" in p:
+            self.weight = float(p["weight"])
+
+    def get_params(self) -> dict:
+        return dict(
+            sphere_collider=dict(center=np.asarray(self.center).tolist(), radius=float(self.radius)),
+            min_clearance=float(self.min_clearance),
+            segment_radius=float(self.segment_radius),
+            weight=float(self.weight),
+        )
+
+    # core -------------------------------------------------------------------
+    def _penalty_single(self, cfg, fk_solver) -> jnp.ndarray:
+        fk = fk_solver.compute_fk_from_angles(cfg)  # (N,4,4)
+        heads = fk[:, :3, 3]                        # (N,3)
+        parents = jnp.asarray(fk_solver.parent_list, jnp.int32)  # (N,)
+        seg_mask = (parents >= 0).astype(jnp.float32)            # (N,)
+        safe_parent_indices = jnp.where(parents >= 0, parents, 0)
+        p_head = heads[safe_parent_indices]
+        c_head = heads
+        v = c_head - p_head
+        dot_vv = jnp.sum(v * v, axis=1) + 1e-6
+        eff_rad = self.radius + self.min_clearance + self.segment_radius
+        vc = self.center - p_head
+        t = jnp.clip(jnp.sum(vc * v, axis=1) / dot_vv, 0.0, 1.0)
+        closest = p_head + t[:, None] * v
+        dist = jnp.linalg.norm(self.center - closest, axis=1)
+        penetration = jnp.maximum(0.0, eff_rad - dist)
+        return jnp.sum((penetration ** 2) * seg_mask)
+
+    def __call__(self, X: jnp.ndarray, fk_solver) -> jnp.ndarray:
+        if X.ndim == 1:
+            loss = self._penalty_single(X, fk_solver)
+        else:
+            loss = jnp.mean(jax.vmap(lambda c: self._penalty_single(c, fk_solver))(X))
+        return loss * jnp.float32(self.weight)

requirements.txt

@@ -1,4 +1,4 @@
-jax-ik==0.1.5
+jax-ik==0.1.6
 wandb
 pillow
 scipy