import logging, faulthandler, sys, time, tempfile, os, requests, zipfile, io, gc, threading, psutil, json, configargparse faulthandler.enable() logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger("jax_ik_server") def _excepthook(t, v, tb): logger.exception("Uncaught exception", exc_info=(t, v, tb)) sys.excepthook = _excepthook # Environment (CPU only) os.environ['JAX_PLATFORMS'] = 'cpu' os.environ['CUDA_VISIBLE_DEVICES'] = '' import jax jax.config.update("jax_default_device", "cpu") import numpy as np from fastapi import FastAPI, Request, Response from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from starlette.responses import FileResponse, JSONResponse from jax_ik.helper import deform_mesh, load_mesh_data_from_gltf, load_mesh_data_from_urdf from jax_ik.hand_specification import HandSpecification from jax_ik.smplx_statics import left_arm_bounds_dict, right_arm_bounds_dict, complete_full_body_bounds_dict from jax_ik.ik import InverseKinematicsSolver from jax_ik.objectives import ( BoneZeroRotationObj, CombinedDerivativeObj, DistanceObjTraj, SphereCollisionPenaltyObjTraj, ) def download_and_setup_files(): os.makedirs("files", exist_ok=True) # Pepper pepper_dir = "files/pepper_description-master" if not os.path.isdir(pepper_dir): logger.info("Downloading Pepper model...") try: r = requests.get("https://uni-bielefeld.sciebo.de/s/98Sy5s143XgNntb/download", stream=True); r.raise_for_status() with zipfile.ZipFile(io.BytesIO(r.content)) as z: z.extractall("files/") except Exception as e: logger.warning(f"Pepper download failed: {e}") # SMPLX smplx_file = "files/smplx.glb" if not os.path.isfile(smplx_file): logger.info("Downloading SMPLX model...") try: r = requests.get("https://uni-bielefeld.sciebo.de/s/B5StwQdiR4DW5mc/download"); r.raise_for_status() open(smplx_file, "wb").write(r.content) except Exception as e: logger.warning(f"SMPLX download failed: {e}") class IKServer: def __init__(self, args): self.args = args self.solve_lock = threading.Lock() self.process = psutil.Process(os.getpid()) # Caches self.solver_cache = {} self.urdf_solver_cache = {} self.max_cache_size = 5 self.cache_access_order = [] self.urdf_cache_access_order = [] # Animation buffers self.animation_frames_agent = [] self.animation_frames_urdf = [] # FastAPI self.app = FastAPI() self.app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"] ) os.makedirs("static", exist_ok=True) os.makedirs("files", exist_ok=True) self.app.mount("/static", StaticFiles(directory="static"), name="static") self.app.mount("/files", StaticFiles(directory="files"), name="files") # Init models (once) self._init_agent() self._setup_agent_objectives() self._init_urdf() self._setup_urdf_objectives() # Cleanup tracking self.last_cleanup_time = time.time() self.cleanup_interval = 30 self.agent_solve_counter = 0 # added self.urdf_solve_counter = 0 # added self._register_routes() logger.info("Server ready.") # ---------- Cache ---------- 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 key = order.pop(0) cache.pop(key, None) gc.collect() def _cache_key(self, bones): return tuple(sorted(bones)) def _get_solver(self, bones, is_urdf=False): key = self._cache_key(bones) 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) if is_urdf: solver = InverseKinematicsSolver( model_file=self.urdf_file, controlled_bones=bones, bounds=None, threshold=self.args.threshold, num_steps=self.args.num_steps, compute_sdf=False, ) else: bounds = [] for b in bones: if b in self.bounds_dict: lower, upper = self.bounds_dict[b] bounds.extend(list(zip(lower, upper))) else: bounds.extend([(-90, 90)] * 3) solver = InverseKinematicsSolver( model_file=self.args.gltf_file, controlled_bones=bones, bounds=bounds, threshold=self.args.threshold, num_steps=self.args.num_steps, compute_sdf=False, ) cache[key] = solver order.append(key) return solver # ---------- Initialization ---------- def _init_agent(self): basic = InverseKinematicsSolver( model_file=self.args.gltf_file, controlled_bones=["left_collar"], bounds=[(-90,90)]*3, threshold=self.args.threshold, num_steps=self.args.num_steps, compute_sdf=False, ) self.available_bones = basic.fk_solver.bone_names self.bounds_dict = complete_full_body_bounds_dict if self.args.hand == "left": self.default_controlled_bones = list(left_arm_bounds_dict.keys()) self.default_end_effector = "left_index3" else: self.default_controlled_bones = list(right_arm_bounds_dict.keys()) self.default_end_effector = "right_index3" 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) 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) self.animation_frames_agent = self._frames_from_angles([self.initial_rotations], False) def _init_urdf(self): self.urdf_file = "files/pepper_description-master/urdf/pepper.urdf" basic = InverseKinematicsSolver( model_file=self.urdf_file, controlled_bones=["LShoulder"], bounds=None, threshold=self.args.threshold, num_steps=self.args.num_steps, compute_sdf=False, ) self.urdf_available_bones = basic.fk_solver.bone_names self.urdf_default_controlled_bones = ["LShoulder","LBicep","LElbow","LForeArm","l_wrist"] self.urdf_default_end_effector = "LFinger13_link" 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) 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) self.animation_frames_urdf = self._frames_from_angles([self.urdf_initial_rotations], True) # ---------- Objectives ---------- def _setup_agent_objectives(self): self.distance_obj = DistanceObjTraj( target_points=np.array([0.0,0.2,0.35]), bone_name=self.current_end_effector, use_head=True, weight=1.0, ) self.collision_obj = SphereCollisionPenaltyObjTraj( {"center":[0.1,0.0,0.35],"radius":0.1}, min_clearance=0.0, weight=1.0, ) def _setup_urdf_objectives(self): self.urdf_distance_obj = DistanceObjTraj( target_points=np.array([0.3,0.3,0.35]), bone_name=self.urdf_current_end_effector, use_head=True, weight=1.0, ) self.urdf_collision_obj = SphereCollisionPenaltyObjTraj( {"center":[0.2,0.0,0.35],"radius":0.1}, min_clearance=0.0, weight=1.0, ) # ---------- Configuration ---------- def configure_agent(self, bones, eff): if not bones: bones = self.default_controlled_bones if eff not in self.available_bones: eff = self.default_end_effector changed = bones != self.current_controlled_bones or eff != self.current_end_effector if bones != self.current_controlled_bones: self.current_controlled_bones = bones self.solver = self._get_solver(bones, is_urdf=False) self.initial_rotations = np.zeros(len(self.solver.controlled_bones)*3, dtype=np.float32) self.best_angles = self.initial_rotations.copy() if eff != self.current_end_effector: self.current_end_effector = eff if changed: self._setup_agent_objectives() return {"controlled_bones": self.current_controlled_bones, "end_effector": self.current_end_effector} def configure_urdf(self, bones, eff): if not bones: bones = self.urdf_default_controlled_bones if eff not in self.urdf_available_bones: eff = self.urdf_default_end_effector changed = bones != self.urdf_current_controlled_bones or eff != self.urdf_current_end_effector if bones != self.urdf_current_controlled_bones: self.urdf_current_controlled_bones = bones self.urdf_solver = self._get_solver(bones, is_urdf=True) 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() if eff != self.urdf_current_end_effector: self.urdf_current_end_effector = eff if changed: self._setup_urdf_objectives() return {"controlled_bones": self.urdf_current_controlled_bones, "end_effector": self.urdf_current_end_effector} # ---------- Objectives build ---------- def _build_agent_objectives(self, payload): 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))}) self.collision_obj.update_params({"weight": float(payload.get("collision_weight",1.0))}) subpoints = int(payload.get("subpoints",1)) mandatory, optional = [], [] if payload.get("distance_enabled", True): mandatory.append(self.distance_obj) if payload.get("collision_enabled", False): 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: optional.append(CombinedDerivativeObj(max_order=3, weights=[float(payload.get("derivative_weight",0.05))]*3)) elif not payload.get("bone_zero_enabled", True) and not payload.get("derivative_enabled", True): optional.append(BoneZeroRotationObj(weight=0.01)) # Hand spec hand_shape = payload.get("hand_shape","None") hand_position = payload.get("hand_position","None") params = { "is_pointing": hand_shape=="Pointing", "is_shaping": hand_shape=="Shaping", "is_flat": hand_shape=="Flat", "look_forward": hand_position=="Look Forward", "look_45_up": hand_position=="Look 45° Up", "look_45_down": hand_position=="Look 45° Down", "look_up": hand_position=="Look Up", "look_down": hand_position=="Look Down", "look_45_x_downwards": hand_position=="Look 45° X Downwards", "look_45_x_upwards": hand_position=="Look 45° X Upwards", "look_x_inward": hand_position=="Look X Inward", "look_to_body": hand_position=="Look to Body", "arm_down": hand_position=="Arm Down", "arm_45_down": hand_position=="Arm 45° Down", "arm_flat": hand_position=="Arm Flat", } if any(params.values()): spec = HandSpecification(**params) optional.extend(spec.get_objectives( left_hand=self.args.hand=="left", controlled_bones=self.current_controlled_bones, full_trajectory=subpoints>1, last_position=True, weight=0.5, )) return mandatory, optional, subpoints 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))}) self.urdf_collision_obj.update_params({"weight": float(payload.get("collision_weight",1.0))}) subpoints = int(payload.get("subpoints",1)) 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) 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: optional.append(CombinedDerivativeObj(max_order=3, weights=[float(payload.get("derivative_weight",0.05))]*3)) elif not payload.get("bone_zero_enabled", True) and not payload.get("derivative_enabled", True): optional.append(BoneZeroRotationObj(weight=0.01)) return mandatory, optional, subpoints # ---------- Solving ---------- def _frames_from_angles(self, angles_seq, is_urdf): frames = [] for ang in angles_seq: if is_urdf: verts = deform_mesh(ang, self.urdf_solver.fk_solver, self.urdf_mesh_data) faces = self.urdf_mesh_data["faces"] else: verts = deform_mesh(ang, self.solver.fk_solver, self.mesh_data) faces = self.mesh_data["faces"] frames.append({"vertices": verts.tolist(), "faces": faces.tolist()}) return frames def solve_agent(self, payload): mand, opt, sub = self._build_agent_objectives(payload) start = time.time() best_angles, obj_val, steps = self.solver.solve( initial_rotations=self.initial_rotations, learning_rate=self.args.learning_rate, mandatory_objective_functions=tuple(mand), optional_objective_functions=tuple(opt), ik_points=sub, verbose=False, ) self.best_angles = best_angles[-1].copy() self.initial_rotations = self.best_angles.copy() self.animation_frames_agent = self._frames_from_angles(best_angles, False) print(f"Solved in {time.time()-start:.2f}s over {steps} steps. Obj: {obj_val:.6f}") self.agent_solve_counter += 1 # added return { "solve_time": time.time()-start, "iterations": steps, "objective": obj_val, "frames": len(self.animation_frames_agent), "solve_id": self.agent_solve_counter, # added } def solve_urdf(self, payload): mand, opt, sub = self._build_urdf_objectives(payload) start = time.time() best_angles, obj_val, steps = self.urdf_solver.solve( initial_rotations=self.urdf_initial_rotations, learning_rate=self.args.learning_rate, mandatory_objective_functions=tuple(mand), optional_objective_functions=tuple(opt), ik_points=sub, verbose=False, ) self.urdf_best_angles = best_angles[-1].copy() self.urdf_initial_rotations = self.urdf_best_angles.copy() self.animation_frames_urdf = self._frames_from_angles(best_angles, True) print( f"Solved in {time.time() - start:.2f}s over {steps} steps. Obj: {obj_val:.6f}" ) self.urdf_solve_counter += 1 # added return { "solve_time": time.time()-start, "iterations": steps, "objective": obj_val, "frames": len(self.animation_frames_urdf), "solve_id": self.urdf_solve_counter, # added } # ---------- Housekeeping ---------- def _cleanup(self): now = time.time() if now - self.last_cleanup_time < self.cleanup_interval: return gc.collect() self.last_cleanup_time = now # ---------- API ---------- def _register_routes(self): @self.app.get("/") def index(): return FileResponse("static/index.html") @self.app.get("/threejs_viewer") def legacy(): return FileResponse("static/index.html") @self.app.get("/animation") def animation(request: Request): model = request.query_params.get("model","agent").lower() data = self.animation_frames_urdf if model == "pepper" else self.animation_frames_agent return Response(content=json.dumps(data), media_type="application/json") @self.app.get("/config") def cfg(model: str = "agent"): if model.lower() == "pepper": return { "model":"pepper", "available_bones": self.urdf_available_bones, "selectable_bones": self.urdf_selectable_bones, "default_controlled_bones": self.urdf_current_controlled_bones, "default_end_effector": self.urdf_current_end_effector, "end_effector_choices": self.urdf_available_bones, "hand_shapes": [], "hand_positions": [], "max_subpoints": 20, # added } return { "model":"agent", "available_bones": self.available_bones, "selectable_bones": self.selectable_bones, "default_controlled_bones": self.current_controlled_bones, "default_end_effector": self.current_end_effector, "end_effector_choices": self.available_bones, "hand_shapes":["None","Pointing","Shaping","Flat"], "hand_positions":[ "None","Look Forward","Look 45° Up","Look 45° Down","Look Up","Look Down", "Look 45° X Downwards","Look 45° X Upwards","Look X Inward","Look to Body", "Arm Down","Arm 45° Down","Arm Flat" ], "max_subpoints": 20, # added } @self.app.post("/configure") async def configure(request: Request): payload = await request.json() model = payload.get("model","agent").lower() if model == "pepper": cfg = self.configure_urdf(payload.get("controlled_bones", []), payload.get("end_effector")) else: cfg = self.configure_agent(payload.get("controlled_bones", []), payload.get("end_effector")) return JSONResponse({"status":"ok","config":cfg}) @self.app.post("/solve") async def solve(request: Request): payload = await request.json() model = payload.get("model","agent").lower() return_mode = payload.get("frames_mode", "auto") self._cleanup() with self.solve_lock: try: if model == "pepper": self.configure_urdf(payload.get("controlled_bones", []), payload.get("end_effector", self.urdf_current_end_effector)) result = self.solve_urdf(payload); result["model"]="pepper" if return_mode == "all" or (return_mode == "auto" and payload.get("subpoints",1) > 1): frames = list(self.animation_frames_urdf) # ensure fresh list else: frames = [self.animation_frames_urdf[-1].copy()] if hasattr(self.animation_frames_urdf[-1],'copy') else [dict(self.animation_frames_urdf[-1])] result["frames_data"] = frames else: self.configure_agent(payload.get("controlled_bones", []), payload.get("end_effector", self.current_end_effector)) result = self.solve_agent(payload); result["model"]="agent" if return_mode == "all" or (return_mode == "auto" and payload.get("subpoints",1) > 1): frames = list(self.animation_frames_agent) else: frames = [self.animation_frames_agent[-1].copy()] if hasattr(self.animation_frames_agent[-1],'copy') else [dict(self.animation_frames_agent[-1])] result["frames_data"] = frames return JSONResponse({"status":"ok","result":result}) except Exception as e: logger.exception("Solve failed") return JSONResponse({"status":"error","message":str(e)}, status_code=500) @self.app.get("/health") def health(): return { "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), } @self.app.get("/favicon.ico") def favicon(): return Response(status_code=204) # ---------- Main ---------- def main(): parser = configargparse.ArgumentParser(description="Inverse Kinematics Solver - Three.js Web UI", default_config_files=["config.ini"]) parser.add("--gltf_file", type=str, default="files/smplx.glb") parser.add("--hand", type=str, choices=["left","right"], default="left") parser.add("--threshold", type=float, default=0.0001) parser.add("--num_steps", type=int, default=100) parser.add("--learning_rate", type=float, default=0.2) parser.add("--subpoints", type=int, default=1) parser.add("--api_port", type=int, default=17861) args = parser.parse_args() download_and_setup_files() srv = IKServer(args) import uvicorn uvicorn.run(srv.app, host="0.0.0.0", port=args.api_port) if __name__ == "__main__": main()