diff --git a/docs/source/api_reference/model_catalogs.rst b/docs/source/api_reference/model_catalogs.rst index 94d291cae..61281d1dd 100644 --- a/docs/source/api_reference/model_catalogs.rst +++ b/docs/source/api_reference/model_catalogs.rst @@ -260,3 +260,58 @@ Other dynamics families :maxdepth: 2 rwm + +WorldParticle inference +----------------------- + +WorldParticle owns frozen networks, strict official checkpoint loading and +sample-ordered features, static boundaries and CSR topology. Inputs and outputs +are world-space particle positions and velocities; external force conditioning +is total acceleration, including gravity. Model-frame conversion uses float32, +while neural inference retains the official BF16 arithmetic. Loading and use +must occur on the caller-owned execution stream. + +.. currentmodule:: embodichain_models.dynamics.worldparticle + +.. autosummary:: + + WorldParticleModel + load_model + +.. autoclass:: WorldParticleModel + :members: __call__ + +.. autofunction:: load_model + +The implementation module exposes the same frozen model contract. + +.. currentmodule:: embodichain_models.dynamics.worldparticle.model + +.. autosummary:: + + WorldParticleModel + load_model + +.. autoclass:: WorldParticleModel + :no-index: + +.. autofunction:: load_model + :no-index: + +The NPZ reader applies task stride before the configured frame slice and then +selects initial position and velocity. Optional reference frames are returned +for evaluation and are never used as future solver inputs. + +.. currentmodule:: embodichain_models.dynamics.worldparticle.data + +.. autosummary:: + + load_sample + +.. autofunction:: load_sample + +The model package imports neither simulation nor examples. Its simulator +Provider belongs to ``embodichain.lab.sim.physics.worldparticle`` and uses +DexSim's native particle inference adapter. Model weights and datasets remain +outside the wheel. The retained ToMe CC-BY-NC-4.0 license requires separate +commercial redistribution review. diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 7541dd47f..ef38bc16b 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -1565,6 +1565,27 @@ embodichain.lab.sim.physics.newton NewtonPhysicsBackend +embodichain.lab.sim.physics.worldparticle +--------------------------------------- + +The per-model Provider validates authored particle order, checkpoint timestep, +gravity and supported scene capabilities before allocating frozen weights. +It supports single-world CUDA particle inference and reset, without native +rigid coupling, gradients or CUDA Graph capture. Model-specific boundaries +replace native collision geometry; SimulationManager callers explicitly set +``enable_default_ground=False`` and pass a ``ProviderNewtonSolverCfg`` object. +Built-in and default solver selection remain unchanged. + +.. currentmodule:: embodichain.lab.sim.physics.worldparticle + +.. autosummary:: + + WorldParticleProvider + +.. autoclass:: WorldParticleProvider + :members: register_custom_attributes, create + + embodichain.lab.sim.motion.planners.base_planner ------------------------------------------------ diff --git a/embodichain/lab/sim/cfg/physics.py b/embodichain/lab/sim/cfg/physics.py index 314eca508..80e6f8b49 100644 --- a/embodichain/lab/sim/cfg/physics.py +++ b/embodichain/lab/sim/cfg/physics.py @@ -373,15 +373,8 @@ def to_dexsim_cfg( ) -> NewtonCfg: """Convert this config to ``dexsim.engine.newton_physics.NewtonCfg``.""" from dexsim.engine.newton_physics import ( - AutoSolverCfg, - DexUniSolverCfg, - FeatherstoneSolverCfg, - MJWarpSolverCfg, NewtonCfg, NewtonCollisionPipelineCfg as DexsimNewtonCollisionPipelineCfg, - SemiImplicitSolverCfg, - VBDSolverCfg, - XPBDSolverCfg, ) torch_device = ( @@ -393,20 +386,34 @@ def to_dexsim_cfg( else str(torch_device) ) - solver_cfg_map: dict[str, type] = { - "auto": AutoSolverCfg, - "dexuni": DexUniSolverCfg, - "mujoco_warp": MJWarpSolverCfg, - "xpbd": XPBDSolverCfg, - "semi_implicit": SemiImplicitSolverCfg, - "featherstone": FeatherstoneSolverCfg, - "vbd": VBDSolverCfg, - } - solver_cfg = _newton_solver_cfg_to_dexsim( - solver_cfg=self.solver_cfg, - solver_cfg_map=solver_cfg_map, - ) + if self.solver_cfg is not None and not isinstance(self.solver_cfg, Mapping): + # Explicit SDK configs (including external Providers) already own + # solver selection. Do not resolve unrelated built-in solver APIs. + solver_cfg = _newton_solver_cfg_to_dexsim(self.solver_cfg, {}) + else: + from dexsim.engine.newton_physics import ( + AutoSolverCfg, + DexUniSolverCfg, + FeatherstoneSolverCfg, + MJWarpSolverCfg, + SemiImplicitSolverCfg, + VBDSolverCfg, + XPBDSolverCfg, + ) + solver_cfg_map: dict[str, type] = { + "auto": AutoSolverCfg, + "dexuni": DexUniSolverCfg, + "mujoco_warp": MJWarpSolverCfg, + "xpbd": XPBDSolverCfg, + "semi_implicit": SemiImplicitSolverCfg, + "featherstone": FeatherstoneSolverCfg, + "vbd": VBDSolverCfg, + } + solver_cfg = _newton_solver_cfg_to_dexsim( + solver_cfg=self.solver_cfg, + solver_cfg_map=solver_cfg_map, + ) if self.requires_grad and ( solver_cfg is None or solver_cfg.solver_type != "semi_implicit" ): diff --git a/embodichain/lab/sim/physics/worldparticle.py b/embodichain/lab/sim/physics/worldparticle.py new file mode 100644 index 000000000..d9955f8bf --- /dev/null +++ b/embodichain/lab/sim/physics/worldparticle.py @@ -0,0 +1,122 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""WorldParticle native inference adapter for DexSim's provider lifecycle.""" + +from __future__ import annotations +from typing import Any + +__all__ = ["WorldParticleProvider"] +import math +from dataclasses import replace + +from dexsim.engine.newton_physics.solver_extensions import ( + NewtonSolverCapabilities, + NewtonSolverExtension, +) + + +class WorldParticleProvider: + """Select one of the four local official WorldParticle networks with explicit static scene inputs. + + Options: ``checkpoint`` and ``config`` paths, an ``inputs`` NPZ with raw + sample-ordered raw features/boundaries/CSR topology, and optional + ``model_to_world`` rotation. Cloth uses its authored surface; duck + deformation uses an unconnected-particle carrier, not an SPH model. + One CUDA particle world and static network boundaries are supported. This + provider supplies inference and native reset, not functional gradients, + native contact/force reports, dynamic rigid bodies, or CUDA Graph capture. + Boundary points are model inputs, not automatically generated collisions. + """ + + capabilities = NewtonSolverCapabilities( + devices=("cuda",), + reset=True, + cuda_graph=False, + reduced_coordinates=False, + external_wrenches=False, + supported_control_fields=(), + runtime_controls=False, + supported_particle_kinds=("fluid",), + ) + + def register_custom_attributes(self, builder: Any, options: dict[str, Any]) -> None: + """Declare the authored topology kind before particle materialization.""" + from embodichain_models.dynamics.worldparticle._runtime import load_config + + cfg = load_config(options.get("config")) + kind = "cloth" if cfg["data"] == "cloth_on_ball" else "fluid" + self.capabilities = replace( + type(self).capabilities, supported_particle_kinds=(kind,) + ) + + def create(self, context: Any, options: dict[str, Any]) -> NewtonSolverExtension: + """Validate the scene before allocating network weights on the GPU.""" + import numpy as np + import torch + from embodichain_models.dynamics.worldparticle import _runtime + from embodichain_models.dynamics.worldparticle.model import WorldParticleModel + + load_assets = _runtime.load_assets + + model = context.model + if not context.device.is_cuda: + raise ValueError("WorldParticle requires CUDA") + if model.requires_grad: + raise ValueError( + "WorldParticle currently supports inference only (requires_grad=False)" + ) + if ( + model.body_count + or model.joint_count + or model.shape_count + or model.world_count > 1 + ): + raise ValueError( + "WorldParticle supports one particle world without rigid bodies, joints or native collision shapes" + ) + if not model.particle_count: + raise ValueError("WorldParticle requires a nonempty fluid particle set") + cfg, arrays, rotation = load_assets(options) + if len(arrays[0]) != model.particle_count: + raise ValueError( + "Input features must match the flattened Newton particle order/count" + ) + if not math.isclose(context.dt, cfg["timestep"], rel_tol=1e-7, abs_tol=1e-10): + raise ValueError("Newton substep dt must match the checkpoint timestep") + expected_gravity = rotation @ np.asarray(cfg["gravity"], dtype=np.float32) + gravity = model.gravity.numpy().reshape(-1, 3) + if not np.allclose(gravity, expected_gravity, atol=1e-6): + raise ValueError( + "Newton gravity must equal model_to_world @ checkpoint gravity" + ) + inverse_mass = model.particle_inv_mass.numpy() + if not np.isfinite(inverse_mass).all() or np.any(inverse_mass <= 0): + raise ValueError("Initially fixed/massless particles are not supported") + from newton import ParticleFlags + + if np.any((model.particle_flags.numpy() & int(ParticleFlags.ACTIVE)) == 0): + raise ValueError("Initially inactive particles are not supported") + from dexsim.engine.newton_physics.particle_inference import ( + ParticleInferenceSolver, + ) + + def create_predictor(device): + inputs = tuple(torch.as_tensor(value, device=device) for value in arrays) + network = _runtime.load_network(cfg, options["checkpoint"], device) + return WorldParticleModel(network, inputs, rotation, cfg) + + solver = ParticleInferenceSolver(model, create_predictor, context.dt) + return NewtonSolverExtension(solver, self.capabilities, close=solver.close) diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index e13265265..55c3baee8 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -248,9 +248,11 @@ def __init__( startup_summary: Literal["compact", "full", "off"] = "compact", dexsim_startup_info: bool = False, scene_node_capacity: int | None = None, + enable_default_ground: bool = True, ) -> None: self.startup_summary = startup_summary self.dexsim_startup_info = dexsim_startup_info + self.enable_default_ground = enable_default_ground self.scene_node_capacity = scene_node_capacity self.width = width self.height = height @@ -323,6 +325,13 @@ def __init__( height: int = 1080 """The height of the simulation window.""" + enable_default_ground: bool = True + """Declare the built-in collision/render plane. + + Disable explicitly for solvers whose boundaries are authored model inputs. + Other scene objects and custom boundary visuals remain caller-owned. + """ + headless: bool = False """Whether to run without an automatically opened native window. @@ -753,7 +762,8 @@ def _initialize_simulation( if self.has_native_renderer: self._init_sim_resources() self.set_default_background() - self._declare_spawn_default_plane() + if sim_config.enable_default_ground: + self._declare_spawn_default_plane() if self.has_native_renderer: self.set_default_global_lighting() @@ -981,7 +991,7 @@ def asset_uids(self) -> List[str]: Returns: List[str]: list of all assets uid. """ - uid_list = ["default_plane"] + uid_list = ["default_plane"] if self.sim_config.enable_default_ground else [] uid_list.extend(list(self._lights.keys())) uid_list.extend(list(self._sensors.keys())) uid_list.extend(list(self._robots.keys())) @@ -1735,7 +1745,7 @@ def prepare(self) -> None: self._env = result.get_arena("default") self._arenas = [result.get_arena(name) for name in scene.arena_names] self.__dict__.pop("arena_offsets", None) - if self._default_plane is None: + if self.sim_config.enable_default_ground and self._default_plane is None: self._bind_default_plane(scene.handles("default_plane")[0]) # Runtime readiness belongs to the SimulationManager. Keep this and @@ -2855,14 +2865,28 @@ def add_deformable_object(self, cfg: DeformableObjectCfg) -> DeformableObject: "Newton deformable particle sets currently require a CUDA device." ) solver_type = self._active_newton_solver_type - supported_solvers = {"auto", "dexuni", "xpbd", "semi_implicit", "vbd"} - if solver_type not in supported_solvers: - raise NotImplementedError( - f"Newton solver {solver_type!r} does not support deformable " - "particle sets; select one of 'auto', 'dexuni', 'xpbd', " - "'semi_implicit', or 'vbd'." - ) + # External Providers declare supported particle kinds at Spawn build; + # DexSim validates those capabilities before materializing this object. + supported_solvers = { + "auto", + "dexuni", + "xpbd", + "semi_implicit", + "vbd", + "provider", + } physics_cfg = self.sim_config.physics_cfg + if solver_type not in supported_solvers: + from dexsim.engine.newton_physics import RegisteredNewtonSolverCfg + + # Registered names retain their entry-point identity rather than + # the direct factory's "provider" type. Spawn validates their + # actual particle capabilities during finalization. + if not isinstance(physics_cfg.solver_cfg, RegisteredNewtonSolverCfg): + raise NotImplementedError( + f"Newton solver {solver_type!r} does not support deformable " + "particle sets; select a particle solver or an external Provider." + ) if isinstance(physics_cfg, NewtonPhysicsCfg) and physics_cfg.requires_grad: raise NotImplementedError( "Newton deformable state mutation is unavailable when " diff --git a/examples/learning/worldparticle/README.md b/examples/learning/worldparticle/README.md new file mode 100644 index 000000000..d67fb0e3d --- /dev/null +++ b/examples/learning/worldparticle/README.md @@ -0,0 +1,91 @@ +# WorldParticle + +The model layer lives in `embodichain_models.dynamics.worldparticle`: original +network, strict checkpoint loading, authored features/topology and coordinate +conversion. It imports neither simulation nor examples. The integration layer +is `embodichain.lab.sim.physics.worldparticle.WorldParticleProvider`; the native +solver adapter is DexSim's model-independent `ParticleInferenceSolver`. +Providers are constructed per scene through the existing Newton lifecycle. + +## Prepared environment + +From this repository, the prepared environment is `.venv-worldparticle`. +The normal `.venv` is unchanged. Linux x86_64, glibc >=2.35, Python 3.12, +Torch 2.7.1+cu128/CXX11 ABI and SM120 are the currently verified operator profile. +NumPy 2.2.6 is used to satisfy EmbodiChain/LeRobot; Torchvision matches Torch. + +To reproduce, build the matching DexSim SDK using its normal CMake workflow +(`DexSim`, `pybind`, `pip-jit-package` targets), preserving its cached Python +interpreter. Use a local PEP 440 package version such as +`-DDEXSIM_PACKAGE_VERSION=0.5.1rc1+worldparticle.6367b2` so it satisfies this +checkout's SDK dependency; the unmodified published SDK lacks the new adapter. +Then pass the produced CPython 3.12 wheel explicitly: + +```bash +uv venv --python 3.12 .venv-worldparticle --seed +uv pip install --python .venv-worldparticle/bin/python 'setuptools>=78.1.1' wheel +SDK_WHEEL=/absolute/path/to/dexsim_engine-0.5.1rc1+worldparticle.6367b2-cp312-cp312-manylinux_2_35_x86_64.whl +uv pip install --python .venv-worldparticle/bin/python --no-build-isolation \ + -e '.[worldparticle]' "$SDK_WHEEL" \ + --extra-index-url https://download.pytorch.org/whl/cu128 \ + --index-strategy unsafe-best-match +uv pip check --python .venv-worldparticle/bin/python +``` + +The requested DexSim source checkout and EmbodiChain's existing default/mapping +Newton configuration have an AutoSolver API mismatch. This change adapts only +explicit SDK configuration objects used by the Provider; it does not change +default solver selection or add an AutoSolver-to-DexUni fallback. Use the +explicit Provider configuration for this environment. The normal `.venv` +remains untouched. + +The optional requirements in `pyproject.toml` are the dependency authority. +`runtime-manifest.json`, the Open3D patch and relocation utility record operator +build/reproduction provenance. No environment, weight or native wheel is committed. + +## Resources and demos + +```bash +source .venv-worldparticle/bin/activate +embodichain run-example learning/worldparticle prepare +embodichain run-example learning/worldparticle demo --scene fancy_fluid +embodichain run-example learning/worldparticle demo --scene cloth_on_ball +embodichain run-example learning/worldparticle demo --scene inverse_design +embodichain run-example learning/worldparticle demo --scene deep_fluid +``` + +The four profiles retain the official checkpoints. Cloth starts at configured +frame 10; duck uses stride 5. Deep Fluid uses generated dam-break initial +conditions with the official model and has no reference trajectory. Duck is a +forward inference demo, not inverse optimization. Rendering uses DexSim's native +1920x1080 OfflineRT window and public Spawn materials/geometry. Repeated display +runs reset and infer fresh states. `--once --hold-seconds 0` exits after one run; +`--headless` runs without a window. `--capture-dir DIR` saves native camera images. + +Assets default to `outputs/worldparticle/assets`, predictions to +`outputs/worldparticle/demos`. The explicit prepare command verifies all 19 +file sizes and SHA256 hashes; `--local-files-only` forbids downloading. +The immutable archive remains on the existing internal resource service; +its original DexSim URL is a storage location, not an SDK runtime dependency. + +For application scenes, pass `ProviderNewtonSolverCfg(provider_factory= +WorldParticleProvider, options={"checkpoint": ..., "config": ..., "inputs": ...})` +to `NewtonPhysicsCfg(solver_cfg=..., use_cuda_graph=False, collision_cfg=None)`. +For `SimulationManager`, set `enable_default_ground=False` because neural +boundaries are authored in the input NPZ. The ordinary manager ground default +and built-in solver selection are unchanged. +Alternatively, pass `RegisteredNewtonSolverCfg(solver_type="worldparticle", options=...)` +to the same physics config; its installed entry point belongs to the EmbodiChain +distribution. Use the checkpoint dt/gravity and matching authored +particle order. Only a single CUDA particle world is supported, without native +rigid collision shapes, gradients or CUDA Graph capture. Static boundaries and +features come from model inputs. Changing them requires closing/recreating the +scene; reset restores physical state and does not reload model assets. + +## Licensing + +WorldParticle's root MIT license and provenance are included in the model package. +Its vendored ToMe merging component retains CC-BY-NC-4.0 (`LICENSE.tome`). +Moving this code does not grant commercial redistribution permission. The model +code currently ships with the EmbodiChain wheel; this integration has not been +commercially cleared. Weight/data rights also need independent release review. diff --git a/examples/learning/worldparticle/__init__.py b/examples/learning/worldparticle/__init__.py new file mode 100644 index 000000000..f38657347 --- /dev/null +++ b/examples/learning/worldparticle/__init__.py @@ -0,0 +1,16 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +from __future__ import annotations diff --git a/examples/learning/worldparticle/__main__.py b/examples/learning/worldparticle/__main__.py new file mode 100644 index 000000000..73e8ebf5c --- /dev/null +++ b/examples/learning/worldparticle/__main__.py @@ -0,0 +1,391 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Four local WorldParticle checkpoints, stepped by DexSim through one Provider.""" + +from __future__ import annotations +import argparse +import json +from pathlib import Path +import time + +import numpy as np +import yaml + +from ._demo_assets import PROFILES, add_boundary_visuals +from embodichain_models.dynamics.worldparticle.data import load_sample + + +def _demo(default_scene="fancy_fluid") -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--scene", choices=tuple(PROFILES), default=default_scene) + parser.add_argument( + "--assets", type=Path, default=Path("outputs/worldparticle/assets") + ) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--config", type=Path) + parser.add_argument("--sample", type=Path) + parser.add_argument("--steps", type=int) + parser.add_argument("--width", type=int, default=1920) + parser.add_argument("--height", type=int, default=1080) + parser.add_argument( + "--start-frame", type=int, help="Override config frame_start after task stride" + ) + parser.add_argument("--output", type=Path) + parser.add_argument( + "--headless", action="store_true", help="No window; no pacing or final hold" + ) + parser.add_argument( + "--renderer", choices=("OFFLINERT", "HYBRID", "FASTRT"), default="OFFLINERT" + ) + parser.add_argument( + "--fps", + type=float, + default=15.0, + help="Display pacing only; checkpoint timestep is unchanged", + ) + parser.add_argument( + "--hold-seconds", + type=float, + default=-1.0, + help="After inference: -1 keeps final frame until window closes/Ctrl+C; 0 exits", + ) + parser.add_argument( + "--capture-dir", + type=Path, + help="Save actual native camera frames for diagnosis", + ) + parser.add_argument( + "--once", action="store_true", help="Run once and hold the final frame" + ) + parser.add_argument( + "--show-boundary-cloud", + action="store_true", + help="Show the original boundary sample surfels", + ) + args = parser.parse_args() + profile = PROFILES[args.scene] + checkpoint = ( + args.checkpoint + or args.assets / "checkpoints" / profile["checkpoint"] / "final.ckpt" + ) + config_path = ( + args.config + or args.assets / "checkpoints" / profile["checkpoint"] / "config.yaml" + ) + sample_path = args.sample or args.assets / profile["sample"] + steps = profile["steps"] if args.steps is None else args.steps + if steps < 1 or args.fps < 0 or min(args.width, args.height) < 1: + parser.error("steps/window size must be positive and fps must be nonnegative") + output = ( + args.output + or Path("outputs/worldparticle/demos") / args.scene / "prediction.npz" + ) + for path in (checkpoint, config_path, sample_path): + if not path.is_file(): + parser.error(f"Missing local asset: {path}") + with config_path.open() as stream: + model_cfg = yaml.safe_load(stream) + frame_start = ( + int(model_cfg.get("frame_start", 0)) + if args.start_frame is None + else args.start_frame + ) + frame_end = model_cfg.get("frame_end") + initial, velocity, raw, box, normals, faces, reference = load_sample( + sample_path, profile, frame_start=frame_start, frame_end=frame_end + ) + rotation = ( + np.eye(3, dtype=np.float32) + if profile["up"] == "z" + else np.asarray( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], np.float32 + ) + ) + q, qd = initial @ rotation.T, velocity @ rotation.T + import dexsim + from dexsim.engine.newton_physics import NewtonCfg, ProviderNewtonSolverCfg + from embodichain.lab.sim.physics.worldparticle import WorldParticleProvider + from dexsim.spawn import ( + SceneBuilder, + FluidRenderMode, + GeometryDesc, + MaterialDesc, + LightDesc, + CameraDesc, + ) + + cfg = dexsim.WorldConfig() + cfg.open_windows = not args.headless + cfg.init_ground = False + cfg.use_default_physics = False + cfg.backend = ( + dexsim.types.Backend.OPENGL + if args.renderer == "OFFLINERT" + else dexsim.types.Backend.VULKAN + ) + cfg.win_config.width = args.width + cfg.win_config.height = args.height + cfg.raytrace_config.spp = 4 + rt_mode = ( + dexsim.types.RTRenderMode.OPTIX_DENOISE + if args.renderer == "OFFLINERT" + else dexsim.types.RTRenderMode.RAW + ) + cfg.set_rt_render_modes(window=rt_mode, offscreen=rt_mode) + cfg.renderer = ( + dexsim.types.Renderer.NORENDER + if args.headless + else getattr(dexsim.types.Renderer, args.renderer) + ) + cfg.newton_cfg = NewtonCfg( + device="cuda:0", + dt=model_cfg["timestep"], + num_substeps=1, + gravity=rotation @ np.asarray(model_cfg["gravity"]), + sync_to_dexsim=not args.headless, + collision_pipeline_cfg=None, + solver_cfg=ProviderNewtonSolverCfg( + provider_factory=WorldParticleProvider, + options=dict( + checkpoint=str(checkpoint), + config=str(config_path), + inputs=str(sample_path), + model_to_world=rotation.tolist(), + ), + ), + ) + cfg.newton_cfg.use_cuda_graph = False + world = dexsim.World(cfg) + builder = SceneBuilder(world) + material = MaterialDesc( + name="prediction", + base_color=(0.18, 0.55, 0.95, 1.0), + roughness=0.45, + double_sided=True, + ) + if args.headless: + material = None + if faces is not None: + builder.add_clothbody( + "prediction", + GeometryDesc.mesh(vertices=q, triangles=faces), + particle_masses=raw[:, 0].astype(np.float32), + particle_radius=model_cfg["particle_radius"], + material=material, + ) + else: + # FluidDesc is the public unconnected-particle carrier. Duck's learned + # feature is not interpreted as mass or as a new SPH material model. + mass, density = ( + (float(raw[0, 0]), float(raw[0, 1])) + if args.scene == "fancy_fluid" + else ( + (float(raw[0, 0]), 1000.0) + if args.scene == "deep_fluid" + else (1.0, 1000.0) + ) + ) + if args.scene == "fancy_fluid" and not np.allclose(raw[:, :2], raw[0, :2]): + raise ValueError("Fancy Fluid demo requires uniform mass/density") + particle_material = ( + None + if args.headless + else MaterialDesc( + name="neural_particles", + base_color=(0.18, 0.58, 0.9, 1.0), + roughness=0.55, + metallic=0.05, + ) + ) + builder.add_fluid( + "prediction", + q, + density=density, + particle_spacing=(mass / density) ** (1.0 / 3.0), + particle_radius=model_cfg["particle_radius"], + material=particle_material, + render_radius=model_cfg["particle_radius"] * 1.6, + render_mode=FluidRenderMode.PARTICLES, + ) + if not args.headless: + add_boundary_visuals( + builder, + box, + normals, + rotation, + args.scene, + model_cfg["particle_radius"], + show_cloud=args.show_boundary_cloud, + ) + pose = np.eye(4) + pose[:3, 3] = q.mean(0) + [2.0, -3.0, 4.0] + builder.add_light( + LightDesc( + name="key", + light_type=dexsim.render.LightType.POINT, + pose=pose, + intensity=150.0, + ) + ) + if args.capture_dir and not args.headless: + builder.add_camera( + CameraDesc( + name="viewer", + width=args.width, + height=args.height, + headless=True, + near=0.01, + far=100.0, + fov=45.0, + ) + ) + scene = builder.finalize() + window = camera = handle = None + try: + if not args.headless: + scene.set_environment_light_emission([0.8, 0.87, 1.0]) + scene.set_environment_light_intensity(80) + handle = scene.get_particle_set("prediction") + handle.set_particle_velocities(qd) + if not args.headless: + world.open_window() + window = world.get_windows() + framing = np.concatenate((q, box @ rotation.T)) + span = max(float(np.ptp(framing, axis=0).max()), 0.8) + center = (framing.min(0) + framing.max(0)) / 2 + eye = center + span * np.array([0.9, -1.1, 0.7]) + window.set_look_at(eye=eye, look_at=center, up=np.array([0.0, 0.0, 1.0])) + if args.capture_dir: + camera = scene.get_camera("viewer") + camera.look_at(eye, center) + camera.open_camera() + camera.render() + print( + f"{profile['title']} — live inference; close window or Ctrl+C to exit", + flush=True, + ) + actual_initial = handle.get_particle_positions().numpy() @ rotation + np.testing.assert_allclose(actual_initial, initial, rtol=0, atol=1e-7) + frames, velocities = [actual_initial.copy()], [velocity.copy()] + if not args.headless: + time.sleep(2.0) # Let the initial scene become visible before inference. + started = time.perf_counter() + for step in range(steps): + if not args.headless and not world.is_window_initialized(): + break + before = time.perf_counter() + world.update() + if camera is not None: + camera.render() + if args.capture_dir and step in (0, steps // 2, steps - 1): + for _ in range(4): + camera.render() + from PIL import Image + + args.capture_dir.mkdir(parents=True, exist_ok=True) + Image.fromarray(np.asarray(camera.get_rgb_map())[..., :3]).save( + args.capture_dir / f"frame-{step+1:04d}.png" + ) + frames.append(handle.get_particle_positions().numpy() @ rotation) + velocities.append(handle.get_particle_velocities().numpy() @ rotation) + if not ( + np.isfinite(frames[-1]).all() and np.isfinite(velocities[-1]).all() + ): + raise RuntimeError(f"Nonfinite neural state at step {step+1}") + if step == 0 or (step + 1) % 20 == 0: + print(f"{args.scene}: {step+1}/{steps}", flush=True) + if not args.headless and args.fps: + time.sleep(max(0.0, 1.0 / args.fps - (time.perf_counter() - before))) + elapsed = time.perf_counter() - started + output.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + output, + position=np.stack(frames), + velocity=np.stack(velocities), + surface_faces=np.empty((0, 3), np.int32) if faces is None else faces, + ) + report = dict( + scene=args.scene, + checkpoint=str(checkpoint), + steps=len(frames) - 1, + particles=len(initial), + all_finite=True, + initial_state_only=True, + source_initial_frame=frame_start * profile["stride"], + source_frame_stride=profile["stride"], + window_size=None if args.headless else [args.width, args.height], + reference_trajectory=profile["reference"], + generated_initial_conditions=args.scene == "deep_fluid", + elapsed_seconds=elapsed, + display_included=not args.headless, + ) + if reference is not None: + count = min(len(reference), len(frames)) + report.update( + reference_frames_compared=count, + position_mse=float( + np.mean((np.stack(frames[:count]) - reference[:count]) ** 2) + ), + ) + output.with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n") + print(f"Saved {output}\n{json.dumps(report,indent=2)}", flush=True) + if not args.headless and args.hold_seconds != 0: + deadline = time.monotonic() + args.hold_seconds + replay_step = 0 + while world.is_window_initialized() and ( + args.hold_seconds < 0 or time.monotonic() < deadline + ): + if args.once: + time.sleep(0.05) + continue + if replay_step == 0: + time.sleep(1.0) + scene.reset() + handle.set_particle_velocities(qd) + print("Restarting live inference", flush=True) + before = time.perf_counter() + world.update() + replay_step = (replay_step + 1) % steps + if args.fps: + time.sleep( + max(0.0, 1.0 / args.fps - (time.perf_counter() - before)) + ) + except KeyboardInterrupt: + print("Stopped by user", flush=True) + finally: + if camera is not None: + camera.close_camera() + window = None + handle = None + scene.close() + + +def main() -> None: + import sys + + if len(sys.argv) > 1 and sys.argv[1] == "prepare": + sys.argv.pop(1) + from ._prepare import main as prepare_main + + prepare_main() + else: + if len(sys.argv) > 1 and sys.argv[1] == "demo": + sys.argv.pop(1) + _demo() + + +if __name__ == "__main__": + main() diff --git a/examples/learning/worldparticle/_demo_assets.py b/examples/learning/worldparticle/_demo_assets.py new file mode 100644 index 000000000..a39b298da --- /dev/null +++ b/examples/learning/worldparticle/_demo_assets.py @@ -0,0 +1,142 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Local official samples and display topology; no dynamics or reference replay.""" + +from __future__ import annotations +from pathlib import Path + +import numpy as np + +PROFILES = { + "fancy_fluid": dict( + checkpoint="fancy_fluid", + sample="data/fancy_fluid/66/66.npz", + steps=199, + title="WorldParticle / Fancy Fluid", + up="y", + reference=True, + stride=1, + ), + "cloth_on_ball": dict( + checkpoint="cloth_on_ball", + sample="data/cloth_on_ball/10/10.npz", + steps=80, + title="WorldParticle / Cloth on Ball", + up="y", + reference=True, + stride=1, + ), + "inverse_design": dict( + checkpoint="inverse_design", + sample="data/inverse_design/0/0.npz", + steps=15, + title="WorldParticle / Duck deformation (forward inference)", + up="z", + reference=True, + stride=5, + ), + "deep_fluid": dict( + checkpoint="deep_lagrange_dataset", + sample="data/deep_lagrange_demo/initial.npz", + steps=199, + title="WorldParticle / Deep Fluid (custom initial state)", + up="y", + reference=False, + stride=1, + ), +} + + +def add_boundary_visuals( + builder, points, normals, rotation, scene_kind, particle_radius, show_cloud=False +): + """Add display-only boundary geometry through Spawn; no collision shapes.""" + from dexsim.spawn import ObjectDesc, RenderDesc, GeometryDesc, MaterialDesc + + material = MaterialDesc( + name="neural_boundary", + base_color=(0.45, 0.49, 0.55, 1.0), + roughness=0.7, + double_sided=True, + ) + if scene_kind == "cloth_on_ball": + floor = np.abs(points[:, 1] - points[:, 1].min()) < 1e-5 + sphere = points[~floor].astype(np.float64) + fit = np.linalg.lstsq( + np.c_[2 * sphere, np.ones(len(sphere))], + np.sum(sphere**2, axis=1), + rcond=None, + )[0] + radii = np.linalg.norm(sphere - fit[:3], axis=1) + radius = radii.mean() + if np.max(np.abs(radii - radius)) < max(1e-5, radius * 0.001): + pose = np.eye(4, dtype=np.float32) + pose[:3, 3] = fit[:3] @ rotation.T + builder.add_sphere( + "boundary_ball", float(radius), pose=pose, material=material + ) + points, normals = points[floor], normals[floor] + if not show_cloud: + world_points = points @ rotation.T + low, high = world_points.min(0), world_points.max(0) + pose = np.eye(4, dtype=np.float32) + pose[:3, 3] = [(low[0] + high[0]) / 2, (low[1] + high[1]) / 2, low[2] - 0.015] + builder.add_cube( + "display_floor", + (max(high[0] - low[0], 0.5), max(high[1] - low[1], 0.5), 0.03), + pose=pose, + material=material, + ) + return + # Sparse surfels visualize the exact input boundary cloud. They are not + # numerical collision geometry and do not replace the full network input. + stride = max(1, int(np.ceil(len(points) / 4000))) + p = points[::stride] @ rotation.T + n = normals[::stride] @ rotation.T + n /= np.maximum(np.linalg.norm(n, axis=1, keepdims=True), 1e-8) + axis = np.tile([0.0, 0.0, 1.0], (len(p), 1)) + axis[np.abs(n[:, 2]) > 0.9] = [1.0, 0.0, 0.0] + u = np.cross(n, axis) + u /= np.maximum(np.linalg.norm(u, axis=1, keepdims=True), 1e-8) + v = np.cross(n, u) + radius = max(float(particle_radius) * 2, 0.008) + vertices = np.stack( + ( + p - radius * u - radius * v, + p + radius * u - radius * v, + p + radius * u + radius * v, + p - radius * u + radius * v, + ), + axis=1, + ).reshape(-1, 3) + indices = ( + np.arange(len(p))[:, None, None] * 4 + np.array([[0, 1, 2], [0, 2, 3]]) + ).reshape(-1, 3) + builder.add_object( + ObjectDesc( + name="boundary_cloud", + physics=None, + renders=[ + RenderDesc.from_geometry( + GeometryDesc.mesh( + vertices=vertices.astype(np.float32), + triangles=indices.astype(np.int32), + ), + material=material, + ) + ], + ) + ) diff --git a/examples/learning/worldparticle/_prepare.py b/examples/learning/worldparticle/_prepare.py new file mode 100644 index 000000000..87cc429f0 --- /dev/null +++ b/examples/learning/worldparticle/_prepare.py @@ -0,0 +1,92 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Explicit resource preparation with pinned archive and per-file identities.""" + +from __future__ import annotations +import argparse +import hashlib +import json +from pathlib import Path +import shutil +import tempfile +from urllib.request import urlopen +import zipfile + +_ARCHIVE = "worldparticle-demo-assets-2eab8d200762.zip" +_SHA256 = "2eab8d200762df1dcab257be6ca5499dd28223d5a91f8e4ce8b3bec3fcf67a7b" +_URL = "http://192.168.3.120/CoreEngine/DexSim/resources/" + _ARCHIVE + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def prepare(root: Path, *, local_files_only: bool = False) -> None: + manifest = json.loads(Path(__file__).with_name("assets-manifest.json").read_text()) + + def verify() -> bool: + return all( + (root / item["path"]).is_file() + and (root / item["path"]).stat().st_size == item["size"] + and _sha256(root / item["path"]) == item["sha256"] + for item in manifest["files"] + ) + + if verify(): + print(f"Verified {len(manifest['files'])} WorldParticle assets: {root}") + return + if local_files_only: + raise ValueError(f"Missing or corrupt WorldParticle assets: {root}") + root.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=root.parent) as directory: + archive = Path(directory) / _ARCHIVE + with urlopen(_URL, timeout=60) as response, archive.open("wb") as output: + shutil.copyfileobj(response, output) + if _sha256(archive) != _SHA256: + raise ValueError("WorldParticle archive SHA256 mismatch") + with zipfile.ZipFile(archive) as bundle: + # Extract only committed, verified file names, never archive paths. + for item in manifest["files"]: + relative = Path(item["path"]) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError("Invalid resource manifest path") + target = root / relative + target.parent.mkdir(parents=True, exist_ok=True) + with ( + bundle.open("worldparticle/" + relative.as_posix()) as source, + target.open("wb") as output, + ): + shutil.copyfileobj(source, output) + if not verify(): + raise ValueError("Extracted WorldParticle file SHA256 mismatch") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--assets", type=Path, default=Path("outputs/worldparticle/assets") + ) + parser.add_argument("--local-files-only", action="store_true") + args = parser.parse_args() + prepare(args.assets, local_files_only=args.local_files_only) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/worldparticle/assets-manifest.json b/examples/learning/worldparticle/assets-manifest.json new file mode 100644 index 000000000..d3f3bff6c --- /dev/null +++ b/examples/learning/worldparticle/assets-manifest.json @@ -0,0 +1,101 @@ +{ + "schema_version": 1, + "resource_directory": "worldparticle", + "files": [ + { + "path": "checkpoints/cloth_on_ball/config.yaml", + "sha256": "146fd67e05f0239f0f9fe0735622ff53070a54935d36f683282fce14160d17bb", + "size": 2266 + }, + { + "path": "checkpoints/cloth_on_ball/final.ckpt", + "sha256": "b39d2d691c1a2be8eb2a990b8e7a06c43ad9b2519073de3855fd0aae4f4bdc61", + "size": 1881230368 + }, + { + "path": "checkpoints/deep_lagrange_dataset/config.yaml", + "sha256": "22b9fe28d50df5ac9b736e672119cd87291b4aeec46c572dbf969d8cfb78500d", + "size": 1850 + }, + { + "path": "checkpoints/deep_lagrange_dataset/final.ckpt", + "sha256": "2183963f89ab8c45d7a47c9f0fe86cae42df3d28d9ec76b3ad58c8827f9fa63e", + "size": 1551226938 + }, + { + "path": "checkpoints/fancy_fluid/config.yaml", + "sha256": "f63b3cc65ecca2df1334e6b465097452cfa85ed23d0b64b0e879757e3ea9be41", + "size": 2156 + }, + { + "path": "checkpoints/fancy_fluid/final.ckpt", + "sha256": "1dce523ef994b6ca80d3cd9de7572de117c1c87ba36d51bf5fcf913be24938d9", + "size": 1880632297 + }, + { + "path": "checkpoints/inverse_design/config.yaml", + "sha256": "cfa44b462ccc8f6a002dc52f5fe3e54dea17483bbb013a84e398c69054321a57", + "size": 2248 + }, + { + "path": "checkpoints/inverse_design/final.ckpt", + "sha256": "7a958a9c6467556208e8bb86f3881f98c64b420faba3a166536177f18403223f", + "size": 1880631328 + }, + { + "path": "cloth-demo-manifest.json", + "sha256": "4b1dd3b64b6fa847b3a1ba213315356c6b7db00533e8ef02fc95deaed5229a4c", + "size": 338 + }, + { + "path": "data/cloth_on_ball/10/10.npz", + "sha256": "e8ff07c0ebd007a0127a6d3242438dd1d824768342cbdec1504de4568b3c1f7c", + "size": 12951512 + }, + { + "path": "data/deep_lagrange_demo/initial.npz", + "sha256": "baa68f4202771a8119320ddb6a64496e25854f8140b5860b28da5d4dfd41701b", + "size": 27262 + }, + { + "path": "data/fancy_fluid/66/66.npz", + "sha256": "9da7c229f89ea892ecee00453e4d3fcba605efa9e6e9ea0bb3d351c7b85cf7bd", + "size": 93993402 + }, + { + "path": "data/inverse_design/0/0.npz", + "sha256": "2f03b6521b6ed82feff316f2d832195cb029fef1031a60bd9382d806bba024d6", + "size": 3714858 + }, + { + "path": "deep-fluid-demo.json", + "sha256": "512875381452861e02e68dbc772e50d5d2e1f040c4be800c856fac30c4a9a5c2", + "size": 1016 + }, + { + "path": "demo-manifest.json", + "sha256": "42f03c5246771b2e390536f20e010aebcacba47892ab4dea9164bb564bb1fd80", + "size": 618 + }, + { + "path": "inverse-demo-manifest.json", + "sha256": "cb2109b7fb14b282afdffdacec72a681a42f345248c9f0799aaa1e439fc16eff", + "size": 757 + }, + { + "path": "weights-manifest.json", + "sha256": "f4ea1e38a06d772ba5b25904d552fd18908f2a1a8f0b64d10d234b592b86b214", + "size": 1920 + }, + { + "path": "data/inverse_design/0/case_input.json", + "sha256": "7e5cc4e424e995d1747cc37b68e1df5f47786d4917329052ffa0a629fcf7521f", + "size": 2243 + }, + { + "path": "data/inverse_design/0/config.json", + "sha256": "54b11940c7bca9090f459f338606bf4666bebbf63b116f5c104827f8b8a6de24", + "size": 7014 + } + ] +} diff --git a/examples/learning/worldparticle/open3d-build.patch b/examples/learning/worldparticle/open3d-build.patch new file mode 100644 index 000000000..0946a2c90 --- /dev/null +++ b/examples/learning/worldparticle/open3d-build.patch @@ -0,0 +1,31 @@ +diff --git a/3rdparty/cmake/FindPytorch.cmake b/3rdparty/cmake/FindPytorch.cmake +index 08bf18b..cef4fd7 100644 +--- a/3rdparty/cmake/FindPytorch.cmake ++++ b/3rdparty/cmake/FindPytorch.cmake +@@ -17,9 +17,9 @@ + # "80-real" to "8.0" and "80" to "8.0+PTX": + macro(translate_arch_string input output) + if("${input}" MATCHES "[0-9]+-real") +- string(REGEX REPLACE "([1-9])([0-9])-real" "\\1.\\2" version "${input}") ++ string(REGEX REPLACE "([1-9][0-9]*)([0-9])-real" "\\1.\\2" version "${input}") + elseif("${input}" MATCHES "([0-9]+)") +- string(REGEX REPLACE "([1-9])([0-9])" "\\1.\\2+PTX" version "${input}") ++ string(REGEX REPLACE "([1-9][0-9]*)([0-9])" "\\1.\\2+PTX" version "${input}") + elseif("${input}" STREQUAL "native") + set(version "Auto") + else() +diff --git a/3rdparty/glfw/glfw.cmake b/3rdparty/glfw/glfw.cmake +index b9573a9..3fd4b7a 100644 +--- a/3rdparty/glfw/glfw.cmake ++++ b/3rdparty/glfw/glfw.cmake +@@ -15,6 +15,10 @@ ExternalProject_Add( + -DGLFW_BUILD_EXAMPLES=OFF + -DGLFW_BUILD_TESTS=OFF + -DGLFW_BUILD_DOCS=OFF ++ -DGLFW_BUILD_WAYLAND=OFF ++ -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} ++ -DCMAKE_INCLUDE_PATH=${CMAKE_INCLUDE_PATH} ++ -DCMAKE_LIBRARY_PATH=${CMAKE_LIBRARY_PATH} + BUILD_BYPRODUCTS + /${Open3D_INSTALL_LIB_DIR}/${CMAKE_STATIC_LIBRARY_PREFIX}${GLFW_LIB_NAME}${CMAKE_STATIC_LIBRARY_SUFFIX} + ) diff --git a/examples/learning/worldparticle/relocate_open3d_wheel.py b/examples/learning/worldparticle/relocate_open3d_wheel.py new file mode 100644 index 000000000..36eb444a9 --- /dev/null +++ b/examples/learning/worldparticle/relocate_open3d_wheel.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +"""Remove build-host search paths and assign a new immutable wheel build tag.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +import re +import shutil +import subprocess +import sys +import tempfile + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("wheel", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--build-tag", required=True) + parser.add_argument("--patchelf", default="patchelf") + args = parser.parse_args() + if not re.fullmatch(r"[0-9][A-Za-z0-9_]*", args.build_tag): + parser.error( + "build tag must start with a digit and contain only letters/digits/underscores" + ) + if not args.wheel.name.startswith("open3d-") or not args.wheel.is_file(): + parser.error("input must be an existing Open3D wheel") + patchelf = shutil.which(args.patchelf) + if not patchelf: + parser.error( + "patchelf is required; select the project's worldparticle-build tooling" + ) + parts = args.wheel.stem.split("-") + name = "-".join(parts[:2] + [args.build_tag] + parts[-3:]) + ".whl" + args.output.mkdir(parents=True, exist_ok=True) + target = args.output.resolve() / name + if target.exists(): + parser.error("refusing to overwrite an existing immutable wheel") + records = [] + with tempfile.TemporaryDirectory(prefix="dexsim-open3d-wheel-") as directory: + root = Path(directory) + subprocess.run( + [ + sys.executable, + "-m", + "wheel", + "unpack", + str(args.wheel.resolve()), + "-d", + str(root), + ], + check=True, + ) + (unpacked,) = root.iterdir() + package = unpacked / "open3d" + for path in sorted(package.rglob("*")): + if not path.is_file(): + continue + with path.open("rb") as stream: + if stream.read(4) != b"\x7fELF": + continue + relatives = [ + os.path.relpath(package, path.parent), + os.path.relpath(unpacked / "torch/lib", path.parent), + os.path.relpath(unpacked / "nvidia/cuda_runtime/lib", path.parent), + ] + runpath = ":".join( + "$ORIGIN" if relative == "." else "$ORIGIN/" + relative + for relative in relatives + ) + subprocess.run([patchelf, "--set-rpath", runpath, str(path)], check=True) + actual = subprocess.check_output( + [patchelf, "--print-rpath", str(path)], text=True + ).strip() + if actual != runpath: + raise RuntimeError("ELF relocation verification failed") + records.append({"path": str(path.relative_to(unpacked)), "runpath": actual}) + if not records: + raise RuntimeError("No ELF libraries found in the Open3D wheel") + subprocess.run( + [ + sys.executable, + "-m", + "wheel", + "pack", + str(unpacked), + "-d", + str(args.output.resolve()), + "--build-number", + args.build_tag, + ], + check=True, + ) + digest = hashlib.sha256() + with target.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + digest = digest.hexdigest() + report = { + "filename": target.name, + "sha256": digest, + "size": target.stat().st_size, + "elf_libraries": records, + } + target.with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps({key: report[key] for key in ["filename", "sha256", "size"]})) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/learning/worldparticle/runtime-manifest.json b/examples/learning/worldparticle/runtime-manifest.json new file mode 100644 index 000000000..97b3761f4 --- /dev/null +++ b/examples/learning/worldparticle/runtime-manifest.json @@ -0,0 +1,51 @@ +{ + "schema_version": 1, + "platform": "linux-x86_64", + "python": "3.12", + "minimum_glibc": "2.35", + "torch": "2.7.1+cu128", + "cxx11_abi": true, + "cuda": "12.8", + "cuda_architectures": [ + "sm_120" + ], + "open3d_source": "https://github.com/isl-org/Open3D", + "open3d_revision": "1e7b17438687a0b0c1e5a7187321ac7044afe275", + "open3d_patch": "open3d-build.patch", + "wheels": [ + { + "filename": "open3d-0.19.0+1e7b174-1torch271cu128sm120-cp312-cp312-manylinux_2_35_x86_64.whl", + "url": "http://192.168.3.43:8080/packages/open3d-0.19.0%2B1e7b174-1torch271cu128sm120-cp312-cp312-manylinux_2_35_x86_64.whl", + "sha256": "209bcee758d25953bcc4cab284ad4fb17f5aace0a3e1e3c296750511bf4257ad", + "size": 443543617 + }, + { + "filename": "flash_attn-2.8.3+cu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", + "url": "http://192.168.3.43:8080/packages/flash_attn-2.8.3%2Bcu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", + "sha256": "cafd4f708725e1be95495a86cbf418cfada08e707c8a797def61488a2a3c4f72", + "size": 256028810 + } + ], + "open3d_patch_sha256": "853748f047ba02fa256c4aa0179c9fa18a67173036be20d8d758a65e5c991a1e", + "numpy": "2.2.6", + "relocation": { + "build_tag": "1torch271cu128sm120", + "input_sha256": "b054609fddb72b97eeb324c7370e94b244b4f78f82552cea859b041b150d00ec", + "elf_libraries": [ + { + "path": "open3d/cuda/open3d_torch_ops.so", + "runpath": "$ORIGIN/..:$ORIGIN/../../torch/lib:$ORIGIN/../../nvidia/cuda_runtime/lib" + }, + { + "path": "open3d/cuda/pybind.cpython-312-x86_64-linux-gnu.so", + "runpath": "$ORIGIN/..:$ORIGIN/../../torch/lib:$ORIGIN/../../nvidia/cuda_runtime/lib" + }, + { + "path": "open3d/libtbb.so.12", + "runpath": "$ORIGIN:$ORIGIN/../torch/lib:$ORIGIN/../nvidia/cuda_runtime/lib" + } + ] + }, + "torchvision": "0.22.1+cu128", + "integration": "EmbodiChain; NumPy2 compatibility requires the recorded end-to-end validation" +} diff --git a/extensions/embodichain_models/README.md b/extensions/embodichain_models/README.md index 430238a1f..5ec963748 100644 --- a/extensions/embodichain_models/README.md +++ b/extensions/embodichain_models/README.md @@ -33,6 +33,7 @@ Use the following canonical packages for the four current model directions: | Model family | Public package | Landing status | | --- | --- | --- | | Ordinary state/visual BC | `embodichain_models.policy.bc` | Target for the BC follow-up; not created yet. | +| WorldParticle | `embodichain_models.dynamics.worldparticle` | Official frozen particle models; native solver integration owned by `lab/sim/physics/worldparticle.py`. | | NeRD | `embodichain_models.dynamics.nerd` | Implemented in this change. | | Robot world model (RWM) | `embodichain_models.dynamics.rwm` | Implemented. | | RigidFormer | `embodichain_models.dynamics.rigidformer` | Native author architecture, official weights and shared-SL recipe. | diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/LICENSE.tome b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/LICENSE.tome new file mode 100644 index 000000000..dfa6588c7 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/LICENSE.tome @@ -0,0 +1,400 @@ + +Attribution-NonCommercial 4.0 International + +======================================================================= + +Creative Commons Corporation ("Creative Commons") is not a law firm and +does not provide legal services or legal advice. Distribution of +Creative Commons public licenses does not create a lawyer-client or +other relationship. Creative Commons makes its licenses and related +information available on an "as-is" basis. Creative Commons gives no +warranties regarding its licenses, any material licensed under their +terms and conditions, or any related information. Creative Commons +disclaims all liability for damages resulting from their use to the +fullest extent possible. + +Using Creative Commons Public Licenses + +Creative Commons public licenses provide a standard set of terms and +conditions that creators and other rights holders may use to share +original works of authorship and other material subject to copyright +and certain other rights specified in the public license below. The +following considerations are for informational purposes only, are not +exhaustive, and do not form part of our licenses. + + Considerations for licensors: Our public licenses are + intended for use by those authorized to give the public + permission to use material in ways otherwise restricted by + copyright and certain other rights. Our licenses are + irrevocable. Licensors should read and understand the terms + and conditions of the license they choose before applying it. + Licensors should also secure all rights necessary before + applying our licenses so that the public can reuse the + material as expected. Licensors should clearly mark any + material not subject to the license. This includes other CC- + licensed material, or material used under an exception or + limitation to copyright. More considerations for licensors: + wiki.creativecommons.org/Considerations_for_licensors + + Considerations for the public: By using one of our public + licenses, a licensor grants the public permission to use the + licensed material under specified terms and conditions. If + the licensor's permission is not necessary for any reason--for + example, because of any applicable exception or limitation to + copyright--then that use is not regulated by the license. Our + licenses grant only permissions under copyright and certain + other rights that a licensor has authority to grant. Use of + the licensed material may still be restricted for other + reasons, including because others have copyright or other + rights in the material. A licensor may make special requests, + such as asking that all changes be marked or described. + Although not required by our licenses, you are encouraged to + respect those requests where reasonable. More_considerations + for the public: + wiki.creativecommons.org/Considerations_for_licensees + +======================================================================= + +Creative Commons Attribution-NonCommercial 4.0 International Public +License + +By exercising the Licensed Rights (defined below), You accept and agree +to be bound by the terms and conditions of this Creative Commons +Attribution-NonCommercial 4.0 International Public License ("Public +License"). To the extent this Public License may be interpreted as a +contract, You are granted the Licensed Rights in consideration of Your +acceptance of these terms and conditions, and the Licensor grants You +such rights in consideration of benefits the Licensor receives from +making the Licensed Material available under these terms and +conditions. + +Section 1 -- Definitions. + + a. Adapted Material means material subject to Copyright and Similar + Rights that is derived from or based upon the Licensed Material + and in which the Licensed Material is translated, altered, + arranged, transformed, or otherwise modified in a manner requiring + permission under the Copyright and Similar Rights held by the + Licensor. For purposes of this Public License, where the Licensed + Material is a musical work, performance, or sound recording, + Adapted Material is always produced where the Licensed Material is + synched in timed relation with a moving image. + + b. Adapter's License means the license You apply to Your Copyright + and Similar Rights in Your contributions to Adapted Material in + accordance with the terms and conditions of this Public License. + + c. Copyright and Similar Rights means copyright and/or similar rights + closely related to copyright including, without limitation, + performance, broadcast, sound recording, and Sui Generis Database + Rights, without regard to how the rights are labeled or + categorized. For purposes of this Public License, the rights + specified in Section 2(b)(1)-(2) are not Copyright and Similar + Rights. + d. Effective Technological Measures means those measures that, in the + absence of proper authority, may not be circumvented under laws + fulfilling obligations under Article 11 of the WIPO Copyright + Treaty adopted on December 20, 1996, and/or similar international + agreements. + + e. Exceptions and Limitations means fair use, fair dealing, and/or + any other exception or limitation to Copyright and Similar Rights + that applies to Your use of the Licensed Material. + + f. Licensed Material means the artistic or literary work, database, + or other material to which the Licensor applied this Public + License. + + g. Licensed Rights means the rights granted to You subject to the + terms and conditions of this Public License, which are limited to + all Copyright and Similar Rights that apply to Your use of the + Licensed Material and that the Licensor has authority to license. + + h. Licensor means the individual(s) or entity(ies) granting rights + under this Public License. + + i. NonCommercial means not primarily intended for or directed towards + commercial advantage or monetary compensation. For purposes of + this Public License, the exchange of the Licensed Material for + other material subject to Copyright and Similar Rights by digital + file-sharing or similar means is NonCommercial provided there is + no payment of monetary compensation in connection with the + exchange. + + j. Share means to provide material to the public by any means or + process that requires permission under the Licensed Rights, such + as reproduction, public display, public performance, distribution, + dissemination, communication, or importation, and to make material + available to the public including in ways that members of the + public may access the material from a place and at a time + individually chosen by them. + + k. Sui Generis Database Rights means rights other than copyright + resulting from Directive 96/9/EC of the European Parliament and of + the Council of 11 March 1996 on the legal protection of databases, + as amended and/or succeeded, as well as other essentially + equivalent rights anywhere in the world. + + l. You means the individual or entity exercising the Licensed Rights + under this Public License. Your has a corresponding meaning. + +Section 2 -- Scope. + + a. License grant. + + 1. Subject to the terms and conditions of this Public License, + the Licensor hereby grants You a worldwide, royalty-free, + non-sublicensable, non-exclusive, irrevocable license to + exercise the Licensed Rights in the Licensed Material to: + + a. reproduce and Share the Licensed Material, in whole or + in part, for NonCommercial purposes only; and + + b. produce, reproduce, and Share Adapted Material for + NonCommercial purposes only. + + 2. Exceptions and Limitations. For the avoidance of doubt, where + Exceptions and Limitations apply to Your use, this Public + License does not apply, and You do not need to comply with + its terms and conditions. + + 3. Term. The term of this Public License is specified in Section + 6(a). + + 4. Media and formats; technical modifications allowed. The + Licensor authorizes You to exercise the Licensed Rights in + all media and formats whether now known or hereafter created, + and to make technical modifications necessary to do so. The + Licensor waives and/or agrees not to assert any right or + authority to forbid You from making technical modifications + necessary to exercise the Licensed Rights, including + technical modifications necessary to circumvent Effective + Technological Measures. For purposes of this Public License, + simply making modifications authorized by this Section 2(a) + (4) never produces Adapted Material. + + 5. Downstream recipients. + + a. Offer from the Licensor -- Licensed Material. Every + recipient of the Licensed Material automatically + receives an offer from the Licensor to exercise the + Licensed Rights under the terms and conditions of this + Public License. + + b. No downstream restrictions. You may not offer or impose + any additional or different terms or conditions on, or + apply any Effective Technological Measures to, the + Licensed Material if doing so restricts exercise of the + Licensed Rights by any recipient of the Licensed + Material. + + 6. No endorsement. Nothing in this Public License constitutes or + may be construed as permission to assert or imply that You + are, or that Your use of the Licensed Material is, connected + with, or sponsored, endorsed, or granted official status by, + the Licensor or others designated to receive attribution as + provided in Section 3(a)(1)(A)(i). + + b. Other rights. + + 1. Moral rights, such as the right of integrity, are not + licensed under this Public License, nor are publicity, + privacy, and/or other similar personality rights; however, to + the extent possible, the Licensor waives and/or agrees not to + assert any such rights held by the Licensor to the limited + extent necessary to allow You to exercise the Licensed + Rights, but not otherwise. + + 2. Patent and trademark rights are not licensed under this + Public License. + + 3. To the extent possible, the Licensor waives any right to + collect royalties from You for the exercise of the Licensed + Rights, whether directly or through a collecting society + under any voluntary or waivable statutory or compulsory + licensing scheme. In all other cases the Licensor expressly + reserves any right to collect such royalties, including when + the Licensed Material is used other than for NonCommercial + purposes. + +Section 3 -- License Conditions. + +Your exercise of the Licensed Rights is expressly made subject to the +following conditions. + + a. Attribution. + + 1. If You Share the Licensed Material (including in modified + form), You must: + + a. retain the following if it is supplied by the Licensor + with the Licensed Material: + + i. identification of the creator(s) of the Licensed + Material and any others designated to receive + attribution, in any reasonable manner requested by + the Licensor (including by pseudonym if + designated); + + ii. a copyright notice; + + iii. a notice that refers to this Public License; + + iv. a notice that refers to the disclaimer of + warranties; + + v. a URI or hyperlink to the Licensed Material to the + extent reasonably practicable; + + b. indicate if You modified the Licensed Material and + retain an indication of any previous modifications; and + + c. indicate the Licensed Material is licensed under this + Public License, and include the text of, or the URI or + hyperlink to, this Public License. + + 2. You may satisfy the conditions in Section 3(a)(1) in any + reasonable manner based on the medium, means, and context in + which You Share the Licensed Material. For example, it may be + reasonable to satisfy the conditions by providing a URI or + hyperlink to a resource that includes the required + information. + + 3. If requested by the Licensor, You must remove any of the + information required by Section 3(a)(1)(A) to the extent + reasonably practicable. + + 4. If You Share Adapted Material You produce, the Adapter's + License You apply must not prevent recipients of the Adapted + Material from complying with this Public License. + +Section 4 -- Sui Generis Database Rights. + +Where the Licensed Rights include Sui Generis Database Rights that +apply to Your use of the Licensed Material: + + a. for the avoidance of doubt, Section 2(a)(1) grants You the right + to extract, reuse, reproduce, and Share all or a substantial + portion of the contents of the database for NonCommercial purposes + only; + + b. if You include all or a substantial portion of the database + contents in a database in which You have Sui Generis Database + Rights, then the database in which You have Sui Generis Database + Rights (but not its individual contents) is Adapted Material; and + + c. You must comply with the conditions in Section 3(a) if You Share + all or a substantial portion of the contents of the database. + +For the avoidance of doubt, this Section 4 supplements and does not +replace Your obligations under this Public License where the Licensed +Rights include other Copyright and Similar Rights. + +Section 5 -- Disclaimer of Warranties and Limitation of Liability. + + a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE + EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS + AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF + ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, + IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, + WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR + PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, + ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT + KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT + ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. + + b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE + TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, + NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, + INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, + COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR + USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN + ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR + DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR + IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. + + c. The disclaimer of warranties and limitation of liability provided + above shall be interpreted in a manner that, to the extent + possible, most closely approximates an absolute disclaimer and + waiver of all liability. + +Section 6 -- Term and Termination. + + a. This Public License applies for the term of the Copyright and + Similar Rights licensed here. However, if You fail to comply with + this Public License, then Your rights under this Public License + terminate automatically. + + b. Where Your right to use the Licensed Material has terminated under + Section 6(a), it reinstates: + + 1. automatically as of the date the violation is cured, provided + it is cured within 30 days of Your discovery of the + violation; or + + 2. upon express reinstatement by the Licensor. + + For the avoidance of doubt, this Section 6(b) does not affect any + right the Licensor may have to seek remedies for Your violations + of this Public License. + + c. For the avoidance of doubt, the Licensor may also offer the + Licensed Material under separate terms or conditions or stop + distributing the Licensed Material at any time; however, doing so + will not terminate this Public License. + + d. Sections 1, 5, 6, 7, and 8 survive termination of this Public + License. + +Section 7 -- Other Terms and Conditions. + + a. The Licensor shall not be bound by any additional or different + terms or conditions communicated by You unless expressly agreed. + + b. Any arrangements, understandings, or agreements regarding the + Licensed Material not stated herein are separate from and + independent of the terms and conditions of this Public License. + +Section 8 -- Interpretation. + + a. For the avoidance of doubt, this Public License does not, and + shall not be interpreted to, reduce, limit, restrict, or impose + conditions on any use of the Licensed Material that could lawfully + be made without permission under this Public License. + + b. To the extent possible, if any provision of this Public License is + deemed unenforceable, it shall be automatically reformed to the + minimum extent necessary to make it enforceable. If the provision + cannot be reformed, it shall be severed from this Public License + without affecting the enforceability of the remaining terms and + conditions. + + c. No term or condition of this Public License will be waived and no + failure to comply consented to unless expressly agreed to by the + Licensor. + + d. Nothing in this Public License constitutes or may be interpreted + as a limitation upon, or waiver of, any privileges and immunities + that apply to the Licensor or You, including from the legal + processes of any jurisdiction or authority. + +======================================================================= + +Creative Commons is not a party to its public +licenses. Notwithstanding, Creative Commons may elect to apply one of +its public licenses to material it publishes and in those instances +will be considered the “Licensor.” The text of the Creative Commons +public licenses is dedicated to the public domain under the CC0 Public +Domain Dedication. Except for the limited purpose of indicating that +material is shared under a Creative Commons public license or as +otherwise permitted by the Creative Commons policies published at +creativecommons.org/policies, Creative Commons does not authorize the +use of the trademark "Creative Commons" or any other trademark or logo +of Creative Commons without its prior written consent including, +without limitation, in connection with any unauthorized modifications +to any of its public licenses or any other arrangements, +understandings, or agreements concerning use of licensed material. For +the avoidance of doubt, this paragraph does not form part of the +public licenses. + +Creative Commons may be contacted at creativecommons.org. diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/LICENSE.worldparticle b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/LICENSE.worldparticle new file mode 100644 index 000000000..15b604692 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/LICENSE.worldparticle @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 WorldParticle 🌐 + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/PROVENANCE.json b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/PROVENANCE.json new file mode 100644 index 000000000..24bac01ac --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/PROVENANCE.json @@ -0,0 +1,29 @@ +{ + "source": "https://github.com/WorldParticle/WorldParticle", + "upstream_revision": "b8e37b854f2af3fc9bfd7fc150cfc01cba4a092e", + "source_files_sha256": { + "models/particle_network_cross_attn_feat.py": "acb65d5a36a26addcf3b1cc7368665d970ec90ae0799988145435f1391035c50", + "models/cross_attn_feat.py": "60f1cd9b755c08f64cb2d67ea326b60a3ac8f5a4cb940a7dd2c12ce15840ab85", + "models/super_particle_layers.py": "126dfdb567b9da5711f221931b0a8dee0859b47fe562b8803a6fcad272b5a675", + "models/continuous_conv.py": "8a4bcccffe06c82a7750be48129c1bbfb686ec0bd66640d389acfc20bf8d1253", + "models/layers/attention.py": "feb20a84177a9448c0e8abb24aef806ef3f6de569d74808a9c63b495eebc95c7", + "models/layers/merge.py": "e01542cd0fee2c7c6e44ea21ad2beaa9860549acffc6f3a3ff84e6d6e4df20dc", + "models/encodings/rope.py": "709e3209e7c427a2e1a44b76eee624b0d361c808b7656f992b37bafbf3d13aed", + "models/encodings/rms_norm.py": "6b90031c3fe097270180d26b0741e0e739056e3d575a75f3315c057320027b69" + }, + "adaptations": [ + "Package-relative imports; removed sys.path mutation", + "Task predicates extracted without dataset/Lightning dependencies", + "Missing FlashAttention fails explicitly instead of changing the attention backend", + "YAPF whitespace formatting; no numerical algorithm changes" + ], + "third_party": { + "ToMe": { + "source": "https://github.com/facebookresearch/ToMe", + "license": "CC-BY-NC-4.0", + "license_file": "LICENSE.tome", + "affected_file": "_network/layers/merge.py", + "notice": "Upstream WorldParticle retains a Meta/ToMe header here. Its license must be considered separately from the WorldParticle root MIT license." + } + } +} diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/__init__.py new file mode 100644 index 000000000..b5a2e5e1e --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/__init__.py @@ -0,0 +1,30 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Lazy WorldParticle inference models, independent of simulation ownership.""" + +from __future__ import annotations +from importlib import import_module +from typing import Any + +__all__ = ["WorldParticleModel", "load_model"] + + +def __getattr__(name: str) -> Any: + if name not in __all__: + raise AttributeError(name) + value = getattr(import_module(".model", __name__), name) + globals()[name] = value + return value diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/_data.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/_data.py new file mode 100644 index 000000000..e5d6b4846 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/_data.py @@ -0,0 +1,71 @@ +"""Task predicates extracted from WorldParticle TrajDataset; no dataset imports.""" + +OBSTACLE_NPZ_DATA = { + "sand", + "elastic", + "cloth_on_ball", + "cloth_on_needle", + "elastic_drag", + "elastic_drop", + "non", + "non_obj", + "non_bunny", + "fancy_fluid", + "cloth_no_self_attention", + "cloth_no_ffn", + "cloth_no_all", + "fancy_flulid_div", + "fancy_fluid_div", + "real_world", + "real_world_override", + "duck_drop", + "hair_ball", +} + +TOPOLOGY_ENABLED_DATA = { + "elastic", + "cloth_on_ball", + "cloth_no_self_attention", + "cloth_no_ffn", + "cloth_no_all", + "cloth_on_needle", + "elastic_drag", + "elastic_drop", + "twist", + "twist_force", + "twist_force_override", + "duck_drop", + "hair_ball", +} + +TOPOLOGY_DEFAULT_SEARCH_RADIUS_BY_DATA = { + "flag": 0.05, + "cloth": 0.05, # backward-compatible alias + "elastic": 0.03, + "cloth_on_needle": 0.005, + "cloth_on_ball": 0.01, + "cloth_no_self_attention": 0.01, + "cloth_no_ffn": 0.01, + "cloth_no_all": 0.01, + "elastic_drag": 0.01, + "elastic_drop": 0.01, + "twist": 0.007, + "twist_force": 0.007, + "twist_force_override": 0.007, + "duck_drop": 0.004, + "hair_ball": 0.008, +} + +OBSTACLE_DATA_TYPES = {"fluid"} | OBSTACLE_NPZ_DATA + + +def uses_topology_neighbors(data): + return data in TOPOLOGY_ENABLED_DATA + + +def get_default_topology_search_radius(data): + return TOPOLOGY_DEFAULT_SEARCH_RADIUS_BY_DATA.get(data, None) + + +def uses_obstacle_features(data): + return data in OBSTACLE_DATA_TYPES diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/continuous_conv.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/continuous_conv.py new file mode 100644 index 000000000..bf5fb557d --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/continuous_conv.py @@ -0,0 +1,400 @@ +import torch + +from open3d.ml.torch.python.layers.neighbor_search import RadiusSearch +from open3d.ml.torch.python import ops +from open3d.ml.torch import classes + + +class FixedRadiusSearch(torch.nn.Module): + """Fixed radius search for 3D point clouds. + + This layer computes the neighbors for a fixed radius on a point cloud. + + Example: + This example shows a neighbor search that returns the indices to the + found neighbors and the distances.:: + + import torch + import open3d.ml.torch as ml3d + + points = torch.randn([20,3]) + queries = torch.randn([10,3]) + radius = 0.8 + + nsearch = ml3d.layers.FixedRadiusSearch(return_distances=True) + ans = nsearch(points, queries, radius) + # returns a tuple of neighbors_index, neighbors_row_splits, and neighbors_distance + + + Arguments: + metric: Either L1, L2 or Linf. Default is L2. + + ignore_query_point: If True the points that coincide with the center of + the search window will be ignored. This excludes the query point if + 'queries' and 'points' are the same point cloud. + + return_distances: If True the distances for each neighbor will be returned. + If False a zero length Tensor will be returned instead. + """ + + def __init__( + self, + metric="L2", + ignore_query_point=False, + return_distances=False, + max_hash_table_size=32 * 2**20, + index_dtype=torch.int32, + **kwargs, + ): + super().__init__() + self.metric = metric + self.ignore_query_point = ignore_query_point + self.return_distances = return_distances + self.max_hash_table_size = max_hash_table_size + assert index_dtype in [torch.int32, torch.int64] + self.index_dtype = index_dtype + + def forward( + self, + points, + queries, + radius, + points_row_splits=None, + queries_row_splits=None, + hash_table_size_factor=1 / 64, + hash_table=None, + return_hash_table=False, + ): + """This function computes the neighbors within a fixed radius for each query point. + + Arguments: + + points: The 3D positions of the input points. It can be a RaggedTensor. + + queries: The 3D positions of the query points. It can be a RaggedTensor. + + radius: A scalar with the neighborhood radius + + points_row_splits: Optional 1D vector with the row splits information + if points is batched. This vector is [0, num_points] if there is + only 1 batch item. + + queries_row_splits: Optional 1D vector with the row splits information + if queries is batched. This vector is [0, num_queries] if there is + only 1 batch item. + + hash_table_size_factor: Scalar. The size of the hash table as fraction + of points. + + hash_table: A precomputed hash table generated with build_spatial_hash_table(). + This input can be used to explicitly force the reuse of a hash table in special + cases and is usually not needed. + Note that the hash table must have been generated with the same 'points' array. + + Returns: + 3 Tensors in the following order + + neighbors_index + The compact list of indices of the neighbors. The corresponding query point + can be inferred from the 'neighbor_count_row_splits' vector. + + neighbors_row_splits + The exclusive prefix sum of the neighbor count for the query points including + the total neighbor count as the last element. The size of this array is the + number of queries + 1. + + neighbors_distance + Stores the distance to each neighbor if 'return_distances' is True. + Note that the distances are squared if metric is L2. + This is a zero length Tensor if 'return_distances' is False. + """ + if isinstance(points, classes.RaggedTensor): + points_row_splits = points.row_splits + points = points.values + if isinstance(queries, classes.RaggedTensor): + queries_row_splits = queries.row_splits + queries = queries.values + + if points_row_splits is None: + points_row_splits = torch.LongTensor([0, points.shape[0]]) + if queries_row_splits is None: + queries_row_splits = torch.LongTensor([0, queries.shape[0]]) + + if hash_table is None: + table = ops.build_spatial_hash_table( + max_hash_table_size=self.max_hash_table_size, + points=points, + radius=radius, + points_row_splits=points_row_splits, + hash_table_size_factor=hash_table_size_factor, + ) + else: + table = hash_table + + result = ops.fixed_radius_search( + ignore_query_point=self.ignore_query_point, + return_distances=self.return_distances, + metric=self.metric, + points=points, + queries=queries, + radius=radius, + points_row_splits=points_row_splits, + queries_row_splits=queries_row_splits, + hash_table_splits=table.hash_table_splits, + hash_table_index=table.hash_table_index, + hash_table_cell_splits=table.hash_table_cell_splits, + index_dtype=self.index_dtype, + ) + + if return_hash_table: + return result, table + + return result + + +class ContinuousConv(torch.nn.Module): + + def __init__( + self, + in_channels, + filters, + kernel_size, + activation=None, + use_bias=True, + kernel_initializer=lambda x: torch.nn.init.uniform_(x, -0.05, 0.05), + bias_initializer=torch.nn.init.zeros_, + align_corners=True, + coordinate_mapping="ball_to_cube_radial", + interpolation="linear", + normalize=False, + offset=None, + use_dense_layer_for_center=False, + dense_kernel_initializer=torch.nn.init.xavier_uniform_, + **kwargs, + ): + super().__init__() + + self.in_channels = in_channels + self.filters = filters + self.kernel_size = kernel_size + self.activation = activation + self.use_bias = use_bias + self.kernel_initializer = kernel_initializer + self.bias_initializer = bias_initializer + self.align_corners = align_corners + self.coordinate_mapping = coordinate_mapping + self.interpolation = interpolation + self.normalize = normalize + self.dense_kernel_initializer = dense_kernel_initializer + + if offset is None: + offset = torch.zeros(size=(3,), dtype=torch.float32) + self.register_buffer("offset", offset) + + self.use_dense_layer_for_center = use_dense_layer_for_center + if self.use_dense_layer_for_center: + self.dense = torch.nn.Linear(self.in_channels, self.filters, bias=False) + self.dense_kernel_initializer(self.dense.weight) + + kernel_shape = (*self.kernel_size, self.in_channels, self.filters) + self.kernel = torch.nn.Parameter( + data=torch.Tensor(*kernel_shape), requires_grad=True + ) + self.kernel_initializer(self.kernel) + + if self.use_bias: + self.bias = torch.nn.Parameter( + data=torch.Tensor(self.filters), requires_grad=True + ) + self.bias_initializer(self.bias) + + def forward( + self, + inp_features, + inp_positions, + out_positions, + extents, + inp_importance, + neighbors_index, + neighbors_row_splits, + neighbors_importance, + ): + # Clone offset to avoid inplace modification issues during backward pass + # The Open3D continuous_conv operation may modify the offset buffer in-place + offset = self.offset.clone() + + # for stats and debugging + num_pairs = neighbors_index.shape[0] + self._avg_neighbors = num_pairs / out_positions.shape[0] + + extents_rank2 = extents + while len(extents_rank2.shape) < 2: + extents_rank2 = torch.unsqueeze(extents_rank2, dim=-1) + + self._conv_values = { + "filters": self.kernel, + "out_positions": out_positions, + "extents": extents_rank2, + "offset": offset, + "inp_positions": inp_positions, + "inp_features": inp_features, + "inp_importance": inp_importance, + "neighbors_index": neighbors_index, + "neighbors_row_splits": neighbors_row_splits, + "neighbors_importance": neighbors_importance, + "align_corners": self.align_corners, + "coordinate_mapping": self.coordinate_mapping, + "interpolation": self.interpolation, + "normalize": self.normalize, + } + + out_features = ops.continuous_conv(**self._conv_values) + + self._conv_output = out_features + + if self.use_dense_layer_for_center: + self._dense_output = self.dense(inp_features) + out_features = out_features + self._dense_output + + if self.use_bias: + out_features += self.bias + if not self.activation is None: + out_features = self.activation(out_features) + + return out_features + + +class ParticleRadiusResearch: + + def __init__( + self, + radius_search_metric="L2", + radius_search_ignore_query_points=True, + window_function=None, + ): + self.radius_search_metric = radius_search_metric + self.radius_search_ignore_query_points = radius_search_ignore_query_points + self.window_function = window_function + + self.fixed_radius_search = FixedRadiusSearch( + metric=self.radius_search_metric, + ignore_query_point=self.radius_search_ignore_query_points, + return_distances=not self.window_function is None, + ) + + self.radius_search = RadiusSearch( + metric=self.radius_search_metric, + ignore_query_point=self.radius_search_ignore_query_points, + return_distances=not self.window_function is None, + normalize_distances=not self.window_function is None, + ) + + def __call__( + self, + inp_positions, + out_positions, + extents, + inp_importance=None, + fixed_radius_search_hash_table=None, + user_neighbors_index=None, + user_neighbors_row_splits=None, + user_neighbors_importance=None, + inp_row_splits=None, + out_row_splits=None, + return_hash_table=False, + ): + + if return_hash_table: + assert ( + len(extents.shape) == 0 + ), "returning hash table only supported for fixed radius search with scalar extent" + + if inp_importance is None: + inp_importance = torch.empty( + (0,), dtype=torch.float32, device=inp_positions.device + ) + + return_distances = not self.window_function is None + + if not user_neighbors_index is None and not user_neighbors_row_splits is None: + + if user_neighbors_importance is None: + neighbors_importance = torch.empty( + (0,), dtype=torch.float32, device=inp_positions.device + ) + else: + neighbors_importance = user_neighbors_importance + + neighbors_index = user_neighbors_index + neighbors_row_splits = user_neighbors_row_splits + + else: + if len(extents.shape) == 0: + radius = 0.5 * extents + if return_hash_table: + self.nns, hash_table = self.fixed_radius_search( + inp_positions, + queries=out_positions, + radius=radius, + points_row_splits=inp_row_splits, + queries_row_splits=out_row_splits, + hash_table=fixed_radius_search_hash_table, + return_hash_table=return_hash_table, + ) + else: + self.nns = self.fixed_radius_search( + inp_positions, + queries=out_positions, + radius=radius, + points_row_splits=inp_row_splits, + queries_row_splits=out_row_splits, + hash_table=fixed_radius_search_hash_table, + return_hash_table=return_hash_table, + ) + if return_distances: + if self.radius_search_metric == "L2": + neighbors_distance_normalized = self.nns.neighbors_distance / ( + radius * radius + ) + else: # L1 + neighbors_distance_normalized = ( + self.nns.neighbors_distance / radius + ) + + elif len(extents.shape) == 1: + radii = 0.5 * extents + self.nns = self.radius_search( + inp_positions, + queries=out_positions, + radii=radii, + points_row_splits=inp_row_splits, + queries_row_splits=out_row_splits, + ) + + else: + raise Exception("extents rank must be 0 or 1") + + if self.window_function is None: + neighbors_importance = torch.empty((0,), dtype=torch.float32) + else: + neighbors_importance = self.window_function( + neighbors_distance_normalized + ) + + neighbors_index = self.nns.neighbors_index + neighbors_row_splits = self.nns.neighbors_row_splits + + if return_hash_table: + return { + "inp_importance": inp_importance, + "neighbors_index": neighbors_index, + "neighbors_row_splits": neighbors_row_splits, + "neighbors_importance": neighbors_importance, + }, hash_table + else: + return { + "inp_importance": inp_importance, + "neighbors_index": neighbors_index, + "neighbors_row_splits": neighbors_row_splits, + "neighbors_importance": neighbors_importance, + } diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/cross_attn_feat.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/cross_attn_feat.py new file mode 100644 index 000000000..45c7ef98e --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/cross_attn_feat.py @@ -0,0 +1,147 @@ +import torch +from torch import nn + +from .encodings.rope import ParticleRotaryEmbedding, freqs_to_cos_sin +from .layers.attention import MultiHeadAttentionWithTOME + + +class CrossAttentionWeightedPoswithFeature(nn.Module): + + def __init__( + self, + dim: int, + filter_extent: float, + n_heads: int = 6, + iters: int = 3, + eps: float = 1e-8, + use_rope: bool = True, # Make RoPE optional + dropout: float = 0.1, + ): + super().__init__() + self.dim = dim + self.filter_extent = filter_extent + self.n_heads = n_heads + self.iters = iters + self.eps = eps + self.nh_f = n_heads + self.use_rope = use_rope + + # RoPE (3D) for inputs & slots - optional + if self.use_rope: + dh_f = dim // n_heads + assert dh_f * n_heads == dim, "dim must be divisible by n_heads" + self.dh_f = dh_f + self.input_rope = ParticleRotaryEmbedding(dim=dh_f) + + # ---------- content head projections (features) ---------- + self.self_attn_with_tome = nn.ModuleList() + for _ in range(self.iters): + self.self_attn_with_tome.append( + MultiHeadAttentionWithTOME( + query_dim=dim, + num_heads=n_heads, + ffn_hidden_dim=dim, + kv_dim=None, # self attention + bias=True, + qk_norm=True, + dropout=dropout, + ) + ) + + def forward( + self, + inputs_feat_batch: torch.Tensor, + input_pos_batch: torch.Tensor, + k_partite: int = -1, + ): + if inputs_feat_batch.dim() != 3 or input_pos_batch.dim() != 3: + raise ValueError( + "inputs_feat_batch and input_pos_batch must have shapes [B, N, C] and [B, N, 3]." + ) + if inputs_feat_batch.shape[:2] != input_pos_batch.shape[:2]: + raise ValueError( + f"Feature/position batch mismatch: {inputs_feat_batch.shape[:2]} vs " + f"{input_pos_batch.shape[:2]}" + ) + + batch_size = inputs_feat_batch.shape[0] + if batch_size > 1: + feature_parts = [] + position_parts = [] + for batch_idx in range(batch_size): + features, positions = self._forward_impl( + inputs_feat_batch=inputs_feat_batch[batch_idx : batch_idx + 1], + input_pos_batch=input_pos_batch[batch_idx : batch_idx + 1], + k_partite=k_partite, + ) + feature_parts.append(features) + position_parts.append(positions) + return torch.cat(feature_parts, dim=0), torch.cat(position_parts, dim=0) + + return self._forward_impl( + inputs_feat_batch=inputs_feat_batch, + input_pos_batch=input_pos_batch, + k_partite=k_partite, + ) + + # ---------------- forward ---------------- + def _forward_impl( + self, + inputs_feat_batch: torch.Tensor, + input_pos_batch: torch.Tensor, + k_partite: int, + ): + """ + inputs_feat : (B, N, D) + input_pos : (B, N, 3) + """ + + # Initialize positions as input positions + pos_batch = input_pos_batch.clone() # (B, N, 3) + + feat = inputs_feat_batch + size_token = None + trace_source = None + + for _iter in range(self.iters): + # Update RoPE embedding with current positions at each iteration + if self.use_rope: + rope_ctx = self.input_rope(pos_batch) + rope_cos, rope_sin = freqs_to_cos_sin(rope_ctx, head_dim=self.dh_f) + + # Calculate number of tokens to remove (with safety check) + current_num_tokens = feat.shape[1] + target_tokens = current_num_tokens // k_partite + num_to_remove_tokens = current_num_tokens - target_tokens + + assert ( + num_to_remove_tokens >= 0 + ), "num_to_remove_tokens should be non-negative" + + feat, size_token, trace_source = self.self_attn_with_tome[_iter]( + query=feat, + num_to_remove_tokens=num_to_remove_tokens, + size_token=size_token, + src_key_padding_mask=None, + rope_cos=rope_cos if self.use_rope else None, + rope_sin=rope_sin if self.use_rope else None, + is_trace_source=True, + k_partite=k_partite, + trace_source=trace_source, + ) # (B, K_new, D) where K_new = current_num_tokens - num_to_remove_tokens + + # Calculate new positions using trace_source matrix + # trace_source: (B, K_new, K_old) - binary matrix indicating which old tokens contribute to new tokens + # Normalize trace_source to get weights (handle case where sum is 0) + trace_weights = trace_source.float() # (B, K_new, K_old) + trace_sum = trace_weights.sum(dim=-1, keepdim=True).clamp( + min=1.0 + ) # (B, K_new, 1) + trace_weights = trace_weights / trace_sum # (B, K_new, K_old) + + # Weighted sum of positions: (B, K_new, K_old) @ (B, K_old, 3) -> (B, K_new, 3) + pos_batch = torch.bmm(trace_weights, input_pos_batch) # (B, K_new, 3) + + slots_pos_batch = pos_batch # Final positions after all iterations + + return feat, slots_pos_batch diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/encodings/rms_norm.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/encodings/rms_norm.py new file mode 100644 index 000000000..e438f4de9 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/encodings/rms_norm.py @@ -0,0 +1,36 @@ +import torch.nn as nn +import torch + + +class RMSNorm(nn.Module): + """ + Root Mean Square Layer Normalization (RMSNorm). + More efficient than LayerNorm, used in modern transformers like LLaMA. + + This implementation handles mixed precision computation by keeping normalization + computations in fp32 for numerical stability. + + Reference: Zhang and Sennrich 2019 - "Root Mean Square Layer Normalization" + """ + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + # Keep weight in fp32 for stability in mixed precision computation + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Store original dtype for output + orig_dtype = x.dtype + + # Convert to fp32 for stable computation + x_fp32 = x.float() + + # RMS normalization: x / RMS(x) * weight + # RMS(x) = sqrt(mean(x^2) + eps) + # Compute in fp32 for numerical stability + rms = torch.sqrt(torch.mean(x_fp32**2, dim=-1, keepdim=True) + self.eps) + x_normed = x_fp32 / rms + + # Apply weight (which is in fp32) and convert back to original dtype + return (self.weight * x_normed).to(orig_dtype) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/encodings/rope.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/encodings/rope.py new file mode 100644 index 000000000..83181bcea --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/encodings/rope.py @@ -0,0 +1,411 @@ +# Modified from https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py + +# Copyright (c) 2021 Phil Wang +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +import math +from math import log +from typing import Optional + +import torch +from torch.nn import Module +from torch.amp import autocast +from torch import nn, einsum, Tensor + +from einops import rearrange, repeat + + +# rotary embedding helper functions +def rotate_half(x): + x = rearrange(x, "... (d r) -> ... d r", r=2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return rearrange(x, "... d r -> ... (d r)") + + +def rotate_half_hf(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@autocast("cuda", enabled=False) +def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2): + dtype = t.dtype + + if t.ndim == 3: + seq_len = t.shape[seq_dim] + freqs = freqs[-seq_len:] + + rot_dim = freqs.shape[-1] + end_index = start_index + rot_dim + + assert ( + rot_dim <= t.shape[-1] + ), f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}" + + # Split t into three parts: left, middle (to be transformed), and right + t_left = t[..., :start_index] + t_middle = t[..., start_index:end_index] + t_right = t[..., end_index:] + + # Apply rotary embeddings without modifying t in place + t_transformed = (t_middle * freqs.cos() * scale) + ( + rotate_half(t_middle) * freqs.sin() * scale + ) + + out = torch.cat((t_left, t_transformed, t_right), dim=-1) + + return out.type(dtype) + + +def freqs_to_cos_sin(freqs, scale=1.0, start_index=0, head_dim=None): + """ + Convert frequencies to cos and sin for rotary embeddings. + + Args: + freqs (torch.Tensor): The frequency tensor of shape (..., n_freqs). + scale (float): The scaling factor for the frequencies. + start_index (int): The starting index of the frequencies. + head_dim (int): The dimension of the head. + """ + if head_dim is not None: + # pad the freqs to match the head_dim + freqs = freqs[..., : freqs.shape[-1] // 2] + left_pad = start_index + right_pad = head_dim // 2 - (left_pad + freqs.shape[-1]) + freqs = torch.cat( + ( + torch.zeros((*freqs.shape[:-1], left_pad), device=freqs.device), + freqs, + torch.zeros((*freqs.shape[:-1], right_pad), device=freqs.device), + ), + dim=-1, + ) + freqs = torch.cat([freqs, freqs], dim=-1) + + cos = freqs.cos() * scale + sin = freqs.sin() * scale + return cos, sin + + +@autocast("cuda", enabled=False) +def apply_rotary_emb_cossin(q, k, cos, sin): + """ + q size: (bsz, n_q_head, seq_len, head_dim) + k size: (bsz, n_kv_head, seq_len, head_dim) + cos size: (bsz, 1, seq_len, head_dim) + sin size: (bsz, 1, seq_len, head_dim) + """ + dtype = q.dtype + rot_dim = cos.shape[-1] + assert ( + rot_dim == q.shape[-1] + ), f"feature dimension {q.shape[-1]} is not equal to rotation dimension {rot_dim}" + + # Apply rotary embeddings without modifying t in place + q = (q * cos) + (rotate_half_hf(q) * sin) + k = (k * cos) + (rotate_half_hf(k) * sin) + + return q.type(dtype), k.type(dtype) + + +@autocast("cuda", enabled=False) +def apply_rotary_emb_one_cossin(one_tensor, cos, sin): + """ + one_tensor size: (bsz, n_head, seq_len, head_dim) + cos size: (bsz, 1, seq_len, head_dim) + sin size: (bsz, 1, seq_len, head_dim) + """ + dtype = one_tensor.dtype + rot_dim = cos.shape[-1] + assert ( + rot_dim == one_tensor.shape[-1] + ), f"feature dimension {one_tensor.shape[-1]} is not equal to rotation dimension {rot_dim}" + + # Apply rotary embeddings without modifying t in place + one_tensor = (one_tensor * cos) + (rotate_half_hf(one_tensor) * sin) + + return one_tensor.type(dtype) + + +class TriangleRotaryEmbedding(Module): + + def __init__( + self, + dim, + hf_format=True, + double_max_freq=False, + ): + """ + TriangleRotaryEmbedding is a class that implements the rotary embedding for the triangle. + + Args: + dim (int): The dimension of the rotary embedding. + hf_format (bool): Whether to use the huggingface RoPE format. + double_max_freq (bool): Whether to double the frequency range. + """ + super().__init__() + + self.hf_format = hf_format + + # log spaced frequencies + max_freq = log(dim // 2 - 1, 2) if not double_max_freq else log(dim - 1, 2) + freqs = 2 ** torch.linspace(0, max_freq, dim // 2) + + self.freqs = nn.Parameter(freqs, requires_grad=False) + + # dummy for device + self.register_buffer("dummy", torch.tensor(0), persistent=False) + + # add apply_rotary_emb as static method + self.apply_rotary_emb = staticmethod(apply_rotary_emb) + + @property + def device(self): + return self.dummy.device + + def get_triangle_freqs(self, pos: Tensor): + # generate all frequencies for all triangles + freqs = self.forward(pos) + freqs = rearrange( + freqs, "batch n_tris n_verts d -> batch 1 n_tris (n_verts d)" + ) # 1 for head dim + + if self.hf_format: + freqs = torch.cat([freqs, freqs], dim=-1) + else: + freqs = repeat(freqs, "... f -> ... (f r)", r=2) + return freqs + + @autocast("cuda", enabled=False) + def forward(self, t: Tensor, seq_len=None, offset=0): + freqs = self.freqs + freqs = einsum("..., f -> ... f", t.type(freqs.dtype), freqs) + + return freqs + + +class ParticleRotaryEmbeddingObsolete(Module): + + def __init__( + self, + dim, + hf_format=True, + double_max_freq=False, + ): + """ + ParticleRotaryEmbedding is a class that implements the rotary embedding for the particle. + + Args: + dim (int): The dimension of the rotary embedding. + hf_format (bool): Whether to use the huggingface RoPE format. + double_max_freq (bool): Whether to double the frequency range. + """ + super().__init__() + + self.hf_format = hf_format + + # log spaced frequencies + max_freq = log(dim // 2 - 1, 2) if not double_max_freq else log(dim - 1, 2) + freqs = 2 ** torch.linspace(0, max_freq, dim // 2) + + self.freqs = nn.Parameter(freqs, requires_grad=False) + + # dummy for device + self.register_buffer("dummy", torch.tensor(0), persistent=False) + + # add apply_rotary_emb as static method + self.apply_rotary_emb = staticmethod(apply_rotary_emb) + + @property + def device(self): + return self.dummy.device + + def get_particle_freqs(self, pos: Tensor): + # generate all frequencies for all particles + freqs = self.forward(pos) + freqs = rearrange( + freqs, "batch n_particles d -> batch 1 n_particles d" + ) # 1 for head dim + + if self.hf_format: + freqs = torch.cat([freqs, freqs], dim=-1) + else: + freqs = repeat(freqs, "... f -> ... (f r)", r=2) + return freqs + + @autocast("cuda", enabled=False) + def forward(self, t: Tensor, seq_len=None, offset=0): + """ + Forward pass for particle positions. + + Args: + t: Tensor of shape (..., 3) containing 3D particle positions (x, y, z) + or (..., 1) for 1D positions + Returns: + Tensor of shape (..., dim//2) containing frequency encodings + """ + freqs = self.freqs + + # Check input dimensionality + input_dim = t.shape[-1] + + if input_dim == 3: # 3D positions (x, y, z) + # For 3D positions, we split frequencies across the 3 spatial dimensions + # Each dimension gets dim//6 frequencies (since dim//2 total, divided by 3) + assert ( + len(freqs) % 3 == 0 + ), f"Number of frequencies {len(freqs)} must be divisible by 3 for 3D positions" + + freqs_per_dim = len(freqs) // 3 + + # Apply frequencies to each spatial dimension separately + # x coordinate modulates first third of frequencies + x_freqs = einsum( + "..., f -> ... f", t[..., 0].type(freqs.dtype), freqs[:freqs_per_dim] + ) + # y coordinate modulates second third of frequencies + y_freqs = einsum( + "..., f -> ... f", + t[..., 1].type(freqs.dtype), + freqs[freqs_per_dim : 2 * freqs_per_dim], + ) + # z coordinate modulates third third of frequencies + z_freqs = einsum( + "..., f -> ... f", + t[..., 2].type(freqs.dtype), + freqs[2 * freqs_per_dim :], + ) + + # Concatenate frequencies from all dimensions + freqs = torch.cat([x_freqs, y_freqs, z_freqs], dim=-1) + + elif input_dim == 1: # 1D positions (scalar per particle) + # For 1D positions, use all frequencies + freqs = einsum("..., f -> ... f", t[..., 0].type(freqs.dtype), freqs) + + else: + raise ValueError( + f"ParticleRotaryEmbedding expects input with last dimension 1 or 3, got {input_dim}" + ) + + return freqs + + +class ParticleRotaryEmbedding(nn.Module): + """ + Rotary Position Embedding for PhysFormer's 3D point clouds. + Handles both state tokens and vertex positions. + """ + + def __init__( + self, + dim: int, + max_freq: float = 10.0, + learnable_freqs: bool = True, + include_time: bool = False, + ): + """ + Args: + dim: Dimension of the embeddings (must be divisible by 2, or 6/8 for 3D) + max_freq: Maximum frequency for sinusoidal embeddings + use_3d: Whether to use 3D positional encoding (x, y, z) + learnable_freqs: Whether to make frequencies learnable + include_time: Whether to include time dimension (for dynamics) + """ + super().__init__() + + self.dim = dim + self.include_time = include_time + + # Determine how many dimensions per coordinate + if include_time: + # 4D: (x, y, z, t) + self.n_coords = 4 + assert ( + dim % (self.n_coords * 2) == 0 + ), f"dim {dim} must be divisible by {self.n_coords * 2}" + self.dim_per_coord = dim // (self.n_coords * 2) + else: + # 3D: (x, y, z) + self.n_coords = 3 + assert ( + dim % (self.n_coords * 2) == 0 + ), f"dim {dim} must be divisible by {self.n_coords * 2}" + self.dim_per_coord = dim // (self.n_coords * 2) + + # Create frequency bands + freqs = torch.logspace(0, math.log10(max_freq), self.dim_per_coord) + + if learnable_freqs: + self.freqs = nn.Parameter(freqs, requires_grad=True) + else: + self.register_buffer("freqs", freqs) + + # For state tokens, we might want learnable position embeddings + self.state_pos_embed = None + + def forward_3d(self, positions: torch.Tensor) -> torch.Tensor: + + # Validate input + n_coords = 4 if self.include_time else 3 + assert positions.shape[-1] == n_coords + + # Shape: positions (..., n_coords) + # Shape: freqs (dim_per_coord,) + + # Compute angles for all coordinates at once + # Reshape positions: (..., n_coords, 1) + # Broadcast with freqs: (..., n_coords, dim_per_coord) + angles = positions.unsqueeze(-1) * self.freqs # (..., n_coords, dim_per_coord) + + # Duplicate each angle for sin/cos pairs + # Shape: (..., n_coords, 2*dim_per_coord) + emb = torch.cat([angles, angles], dim=-1) + + # Reshape to concatenate coordinate embeddings + # From: (..., n_coords, 2*dim_per_coord) + # To: (..., n_coords * 2*dim_per_coord) + rope_emb = emb.flatten(-2, -1) + rope_emb = rope_emb.unsqueeze(1) + + return rope_emb + + def forward( + self, + positions: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Flexible forward function that handles different input types. + + Args: + positions: 3D/4D positions for vertices + batch_size: Batch size for generating state token positions + seq_len: Sequence length for state tokens + + Returns: + RoPE embeddings + """ + if positions is not None: + return self.forward_3d(positions) + else: + raise ValueError("Must provide positions (3D or 4D) for RoPE") diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/layers/attention.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/layers/attention.py new file mode 100644 index 000000000..085f4cfa4 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/layers/attention.py @@ -0,0 +1,1038 @@ +import os +from typing import Literal, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ..encodings.rope import ( + ParticleRotaryEmbedding, + apply_rotary_emb_cossin, + apply_rotary_emb_one_cossin, + apply_rotary_emb, +) + +from ..encodings.rms_norm import RMSNorm + +EPS = 1e-6 + +ATTN = os.environ.get("ATTN_IMPL", "flash_attn") +assert ATTN in ["flash_attn", "sdpa"], "ATTN_IMPL must be either 'flash_attn' or 'sdpa'" +if ATTN == "flash_attn": + try: + from flash_attn import ( + flash_attn_qkvpacked_func, + flash_attn_varlen_qkvpacked_func, + flash_attn_varlen_kvpacked_func, + ) + from flash_attn.bert_padding import pad_input, unpad_input + from einops import rearrange + except ImportError: + raise ImportError( + "WorldParticle requires a FlashAttention build matching PyTorch/CUDA. Install the optional WorldParticle runtime before selecting this solver." + ) + + +class FeedForwardSwiGLU(nn.Module): + + def __init__( + self, + dim: int, + hidden_dim: int, + dropout: float = 0.1, + bias: bool = True, + ): + """ + Feed forward layer with SwiGLU activation. + Args: + dim (int): input dimension + hidden_dim (int): feed forward hidden dim + dropout (float): dropout rate, default 0.1 + """ + super().__init__() + + self.w1 = nn.Linear(dim, hidden_dim, bias=bias) + self.w2 = nn.Linear(hidden_dim, dim, bias=bias) + self.w3 = nn.Linear(dim, hidden_dim, bias=bias) + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + + def forward(self, x): + return self.dropout(self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))) + + +class FeedForwardGeLU(nn.Module): + + def __init__( + self, + dim: int, + hidden_dim: int, + dropout: float = 0.1, + bias: bool = True, + ): + """ + Feed forward layer with GeLU activation. + Args: + dim (int): input dimension + hidden_dim (int): feed forward hidden dim + dropout (float): dropout rate, default 0.1 + """ + super().__init__() + + self.w1 = nn.Linear(dim, hidden_dim, bias=bias) + self.w2 = nn.Linear(hidden_dim, dim, bias=bias) + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + + def forward(self, x): + return self.dropout(self.w2(self.dropout(F.gelu(self.w1(x))))) + + +class MultiHeadAttention(nn.Module): + + def __init__( + self, + query_dim, + num_heads, + kv_dim=None, + bias=True, + qk_norm=False, + norm_type="layer_norm", + ): + super().__init__() + self.apply_rope_cossin = apply_rotary_emb_cossin + + self.num_heads = num_heads + self.is_self_attn = kv_dim is None + kv_dim = query_dim if kv_dim is None else kv_dim + + if self.is_self_attn: + self.in_proj = nn.Linear(query_dim, 3 * query_dim, bias=bias) + else: + self.q_proj = nn.Linear(query_dim, query_dim, bias=bias) + self.k_proj = nn.Linear(kv_dim, query_dim, bias=bias) + self.v_proj = nn.Linear(kv_dim, query_dim, bias=bias) + self.out_proj = nn.Linear(query_dim, query_dim, bias=bias) + + if qk_norm: + if norm_type == "layer_norm": + norm_module = nn.LayerNorm + elif norm_type == "rms_norm": + norm_module = RMSNorm + else: + raise ValueError( + "Unsupported normalization type. Choose from 'layer_norm' and 'rms_norm'." + ) + self.q_norm = norm_module(query_dim, eps=EPS) + self.k_norm = norm_module(query_dim, eps=EPS) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + def forward( + self, + q, + k, + v, + src_key_padding_mask=None, + rope_cos=None, + rope_sin=None, + rope_ctx_cos=None, + rope_ctx_sin=None, + force_sdpa=False, + ): + # src_key_padding_mask: (B, N), key padding mask, things you want to attend to is True + bs, src_len = q.shape[0], q.shape[1] + ctx_len = k.shape[1] + + if self.is_self_attn: + q, k, v = self.in_proj(q).chunk(3, dim=-1) + else: + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # qk normalization + q = self.q_norm(q).type(v.dtype) + k = self.k_norm(k).type(v.dtype) + + q = q.view(bs, src_len, self.num_heads, -1).transpose( + 1, 2 + ) # (bs, num_heads, src_len, head_dim) + k = k.view(bs, ctx_len, self.num_heads, -1).transpose(1, 2) + v = v.view(bs, ctx_len, self.num_heads, -1).transpose(1, 2) + + # apply rope + if rope_cos is not None: + if rope_ctx_cos is None: + q, k = self.apply_rope_cossin(q, k, rope_cos, rope_sin) + else: + q = apply_rotary_emb_one_cossin(q, rope_cos, rope_sin) + k = apply_rotary_emb_one_cossin(k, rope_ctx_cos, rope_ctx_sin) + + if ATTN == "sdpa" or force_sdpa: + # create attention mask + if src_key_padding_mask is not None: + # things you want to attend to is True + assert src_key_padding_mask.shape == ( + bs, + ctx_len, + ), f"expecting key_padding_mask shape of {(bs, ctx_len)}, but got {src_key_padding_mask.shape}" + # but now, things you want to attend to is True + src_key_padding_mask = torch.zeros_like( + src_key_padding_mask, dtype=torch.bool + ).masked_fill_(src_key_padding_mask, True) + attn_mask = ( + src_key_padding_mask.view(bs, 1, 1, ctx_len) + .expand(-1, self.num_heads, -1, -1) + .reshape(bs * self.num_heads, 1, ctx_len) + ) + attn_mask = attn_mask.view(bs, self.num_heads, -1, ctx_len) + else: + attn_mask = None + + attn_output = ( + F.scaled_dot_product_attention( + query=q.type(v.dtype), + key=k.type(v.dtype), + value=v, + attn_mask=attn_mask, + ) + .transpose(1, 2) + .contiguous() + .view(bs, src_len, -1) + ) + elif ATTN == "flash_attn": + # self-attn + if self.is_self_attn: + if src_key_padding_mask is not None: + q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, _ = unpad_input( + q.transpose(1, 2), src_key_padding_mask + ) + k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, _ = unpad_input( + k.transpose(1, 2), src_key_padding_mask + ) + v_unpad, indices_v, cu_seqlens_v, max_seqlen_v, _ = unpad_input( + v.transpose(1, 2), src_key_padding_mask + ) + qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1) + out_unpad = flash_attn_varlen_qkvpacked_func( + qkv_unpad, + cu_seqlens_q, + max_seqlen_q, + ) + attn_output = ( + pad_input(out_unpad, indices_q, bs, src_len) + .contiguous() + .view(bs, src_len, -1) + ) + else: + qkv = torch.stack( + [ + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + ], + dim=2, + ) + attn_output = ( + flash_attn_qkvpacked_func(qkv) + .contiguous() + .view(bs, src_len, -1) + ) + # cross-attn + else: + if src_key_padding_mask is not None: + # With padding mask: use variable-length attention + q_unpad = rearrange(q, "b h s d -> (b s) h d") + cu_seqlens_q = torch.arange( + 0, + (bs + 1) * src_len, + step=src_len, + dtype=torch.int32, + device=q_unpad.device, + ) + max_seqlen_q = src_len + + k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, _ = unpad_input( + k.transpose(1, 2), src_key_padding_mask + ) + v_unpad, _, _, _, _ = unpad_input( + v.transpose(1, 2), src_key_padding_mask + ) + kv_unpad = torch.stack([k_unpad, v_unpad], dim=1) + + out_unpad = flash_attn_varlen_kvpacked_func( + q_unpad, + kv_unpad, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + ) + attn_output = ( + rearrange(out_unpad, "(b s) h d -> b s h d", b=bs) + .contiguous() + .view(bs, src_len, -1) + ) + else: + # Without padding mask: use regular flash attention + # Reshape for flash_attn_kvpacked_func + from flash_attn import flash_attn_kvpacked_func + + # q: [bs, num_heads, src_len, head_dim] -> [bs, src_len, num_heads, head_dim] + q_flash = q.transpose(1, 2) + # k, v: [bs, num_heads, ctx_len, head_dim] -> [bs, ctx_len, num_heads, head_dim] + k_flash = k.transpose(1, 2) + v_flash = v.transpose(1, 2) + + # Stack k and v for kvpacked format + kv = torch.stack( + [k_flash, v_flash], dim=2 + ) # [bs, ctx_len, 2, num_heads, head_dim] + + # Apply flash attention + attn_output = ( + flash_attn_kvpacked_func(q_flash, kv) + .contiguous() + .view(bs, src_len, -1) + ) + else: + raise ValueError( + "Unsupported attention type. Choose from 'flash_attn' and 'sdpa'." + ) + + return self.out_proj(attn_output) + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = ( + x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + ) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view( + B, H // window_size, W // window_size, window_size, window_size, -1 + ) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +SWIN_ATTN_MASK_CACHE = {} + + +def get_swin_attn_mask(H, W, window_size, shift_size, device): + """ + Get the attention mask for Swin Transformer. (Original implementation) + Args: + H (int): height of image + W (int): width of image + window_size (int): window size + shift_size (int): shift size + device (torch.device): device to store the attention mask + Returns: + attn_mask: (num_windows, num_windows) + """ + if (H, W, window_size, shift_size) in SWIN_ATTN_MASK_CACHE: + return SWIN_ATTN_MASK_CACHE[(H, W, window_size, shift_size)] + else: + img_mask = torch.zeros((1, H, W, 1), device=device) # 1 H W 1 + h_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + w_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition( + img_mask, window_size + ) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, window_size * window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = (attn_mask == 0).to( + torch.bool + ) # nW, window_size * window_size, window_size * window_size + SWIN_ATTN_MASK_CACHE[(H, W, window_size, shift_size)] = attn_mask + return attn_mask + + +class SwinSelfAttention(nn.Module): + + def __init__( + self, + dim, + num_heads, + window_size, + shift_size: int = 0, + bias=True, + qk_norm=False, + norm_type="layer_norm", + ): + """ + Args: + dim (int): input dimension + num_heads (int): number of attention heads + window_size (int): window size + shift_size (int): shift size, if None, no shift + bias (bool): whether to use bias, default True + qk_norm (bool): whether to normalize query and key, default False + """ + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + + self.in_proj = nn.Linear(dim, 3 * dim, bias=bias) + self.out_proj = nn.Linear(dim, dim, bias=bias) + + if qk_norm: + if norm_type == "layer_norm": + norm_module = nn.LayerNorm + elif norm_type == "rms_norm": + norm_module = RMSNorm + else: + raise ValueError( + "Unsupported normalization type. Choose from 'layer_norm' and 'rms_norm'." + ) + self.q_norm = norm_module(dim, eps=EPS) + self.k_norm = norm_module(dim, eps=EPS) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + def forward(self, x): + """ + Args: + x: (B, H, W, C) + Returns: + x: (B, H, W, C) + """ + B, H, W, C = x.shape + nW = H * W // self.window_size // self.window_size + + # swin related operations + if self.shift_size > 0: + attn_mask = get_swin_attn_mask( + H, W, self.window_size, self.shift_size, x.device + ) # nW, window_size * window_size, window_size * window_size + attn_mask = attn_mask.repeat(B, 1, 1)[ + :, None + ] # B * nW, 1, window_size * window_size, window_size * window_size + else: + attn_mask = None + + if self.shift_size > 0: + shifted_x = torch.roll( + x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) + ) + else: + shifted_x = x + + # partition windows + x_windows = window_partition( + shifted_x, self.window_size + ) # B * nW, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # B * nW, window_size * window_size, C + + q, k, v = self.in_proj(x_windows).chunk(3, dim=-1) + + # qk normalization + q = self.q_norm(q) + k = self.k_norm(k) + + # (B * nW, num_heads, window_size * window_size, head_dim) + q = q.view( + B * nW, self.window_size * self.window_size, self.num_heads, -1 + ).transpose(1, 2) + k = k.view( + B * nW, self.window_size * self.window_size, self.num_heads, -1 + ).transpose(1, 2) + v = v.view( + B * nW, self.window_size * self.window_size, self.num_heads, -1 + ).transpose(1, 2) + + # apply attention + attn_output = ( + F.scaled_dot_product_attention( + query=q.type(v.dtype), + key=k.type(v.dtype), + value=v, + attn_mask=attn_mask, + ) + .transpose(1, 2) + .contiguous() + .view(B * nW, self.window_size * self.window_size, -1) + ) + + attn_windows = self.out_proj( + attn_output + ) # B * nW, window_size * window_size, C + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + # reverse cyclic shift + shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H W C + if self.shift_size > 0: + x = torch.roll( + shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) + ) + else: + x = shifted_x + x = x.view(B, H, W, C) + return x + + +class AttentionLayer(nn.Module): + + def __init__( + self, + query_dim: int, + num_heads: int, + ffn_hidden_dim: int, + kv_dim: Optional[int] = None, + dropout: float = 0.1, + bias: bool = True, + bias_kv: bool = False, + activation: str = "swiglu", + norm_type: Literal["layer_norm", "rms_norm"] = "layer_norm", + disable_q_norm: bool = False, + disable_kv_norm: bool = False, + qk_norm: bool = False, + add_self_attn: bool = False, + use_swin_attn: bool = False, + window_size: int = 8, + shift_size: int = 0, + ): + """ + Attention layer with feed forward and pre-norm. + Args: + query_dim (int): input dimension + kv_dim (int): key and value dimension, if None, set to query_dim (self-attention) + num_heads (int): number of attention heads + hidden_dim (int): feed forward hidden dim + dropout (float): dropout, default 0.1 + bias (bool): whether to use bias, default True + bias_kv (bool): whether to use bias for key and value, default False + activation (str): activation function, choose from 'gelu' and 'swiglu', default 'swiglu' + norm_type (str): normalization type, choose from 'layer_norm' and 'rms_norm', default 'layer_norm' + disable_q_norm (bool): disable query normalization, default False + disable_kv_norm (bool): disable key and value normalization, default False + qk_norm (bool): whether to apply normalization to query and key, default False + add_self_attn (bool): whether to add self-attention after cross-attention (cross-attn, self-attn, ffn), default False + use_swin_attn (bool): whether to use swin self-attention, default False + window_size (int): window size for swin self-attention, default 8 + shift_size (int): shift size for swin self-attention, default 0 (no shift) + Returns: + torch.Tensor: (B, N, query_dim) + """ + super().__init__() + self.multihead_attn = MultiHeadAttention( + query_dim=query_dim, + num_heads=num_heads, + kv_dim=kv_dim, + bias=bias, + qk_norm=qk_norm, + norm_type=norm_type, + ) + print(f"dropout: {dropout}") + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + + if bias_kv: + raise NotImplementedError("Bias for key and value is not supported for now") + + if norm_type == "layer_norm": + norm_module = nn.LayerNorm + elif norm_type == "rms_norm": + norm_module = RMSNorm + else: + raise ValueError( + "Unsupported normalization type. Choose from 'layer_norm' and 'rms_norm'." + ) + + self.query_norm = ( + norm_module(query_dim, eps=EPS) if not disable_q_norm else nn.Identity() + ) + kv_dim = query_dim if kv_dim is None else kv_dim + if not self.multihead_attn.is_self_attn: + self.kv_norm = ( + norm_module(kv_dim, eps=EPS) if not disable_kv_norm else nn.Identity() + ) + + self.add_self_attn = add_self_attn + self.use_swin_attn = use_swin_attn + if add_self_attn: + if use_swin_attn: + self.self_attn = SwinSelfAttention( + dim=query_dim, + num_heads=num_heads, + window_size=window_size, + shift_size=shift_size, + bias=bias, + qk_norm=qk_norm, + norm_type=norm_type, + ) + else: + self.self_attn = MultiHeadAttention( + query_dim=query_dim, + num_heads=num_heads, + kv_dim=None, + bias=bias, + qk_norm=qk_norm, + norm_type=norm_type, + ) + self.self_attn_norm = ( + norm_module(query_dim, eps=EPS) if not disable_q_norm else nn.Identity() + ) + + if activation == "swiglu": + self.ffn = FeedForwardSwiGLU( + query_dim, hidden_dim=ffn_hidden_dim, dropout=dropout, bias=bias + ) + elif activation == "gelu": + self.ffn = FeedForwardGeLU( + query_dim, hidden_dim=ffn_hidden_dim, dropout=dropout, bias=bias + ) + else: + raise ValueError( + "Unsupported activation function. Choose from 'gelu' and 'swiglu'." + ) + + self.ffn_norm = norm_module(query_dim, eps=EPS) + + def forward( + self, + query, + kv=None, + src_key_padding_mask=None, + rope_cos=None, + rope_sin=None, + rope_ctx_cos=None, + rope_ctx_sin=None, + force_sdpa=False, + patch_h=None, + patch_w=None, + ): + """ + Args: + query (torch.Tensor): (B, N, query_dim) + kv (torch.Tensor): (B, N, kv_dim), key and value, if None, set to query (self-attention) + src_key_padding_mask (torch.Tensor): (B, N), key padding mask, things you want to attend to is True + rope_cos (torch.Tensor): (B, 1, N, head_dim), cosine tensor for RoPE, None if no RoPE is applied + rope_sin (torch.Tensor): (B, 1, N, head_dim), sine tensor for RoPE, None if no RoPE is applied + rope_ctx_cos (torch.Tensor): (B, 1, N, head_dim), cosine tensor for RoPE, None if no RoPE is applied + rope_ctx_sin (torch.Tensor): (B, 1, N, head_dim), sine tensor for RoPE, None if no RoPE is applied + patch_h (int): height of the patch, used for swin self-attention + patch_w (int): width of the patch, used for swin self-attention + Returns: + torch.Tensor: (B, N, query_dim) + """ + + bs = query.shape[0] + + q = self.query_norm(query) + if self.multihead_attn.is_self_attn: + kv = q + else: + kv = self.kv_norm(kv) + + # multihead attention + attn_output = self.dropout( + self.multihead_attn( + q, + kv, + kv, + src_key_padding_mask, + rope_cos, + rope_sin, + rope_ctx_cos, + rope_ctx_sin, + force_sdpa=force_sdpa, + ) + ) + query = query + attn_output + + if self.add_self_attn: + q = self.self_attn_norm(query) + if self.use_swin_attn: + q = q.view(bs, patch_h, patch_w, -1) + self_attn_output = self.self_attn(q) + self_attn_output = self_attn_output.view(bs, patch_h * patch_w, -1) + else: + self_attn_output = self.self_attn( + q, q, q, None, rope_cos, rope_sin, force_sdpa=force_sdpa + ) + query = query + self.dropout(self_attn_output) + + # feed forward + query = query + self.dropout(self.ffn(self.ffn_norm(query))) + return query + + +class CrossAttentionWithRoPE(nn.Module): + """ + Cross-attention that applies RoPE to queries/keys. + + RenderFormer enhancement: QK-normalization for numerical stability. + """ + + def __init__( + self, + d_model: int, + num_heads: int, + dropout: float = 0.1, + use_rope: bool = True, + qk_norm: bool = True, # RenderFormer: QK-normalization + rope_dim: int = None, # Custom RoPE dimension (None = auto-compute) + ): + super().__init__() + assert d_model % num_heads == 0, "d_model must be divisible by num_heads" + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + self.scale = self.head_dim**-0.5 + self.use_rope = bool(use_rope) + self.qk_norm = qk_norm + + # Projections + self.q_proj = nn.Linear(d_model, d_model, bias=True) + self.k_proj = nn.Linear(d_model, d_model, bias=True) + self.v_proj = nn.Linear(d_model, d_model, bias=True) + self.out_proj = nn.Linear(d_model, d_model, bias=True) + self.dropout = nn.Dropout(dropout) + + # QK-normalization: RMSNorm per head (RenderFormer recommendation) + if self.qk_norm: + self.q_norm = RMSNorm(self.head_dim) + self.k_norm = RMSNorm(self.head_dim) + + # 3D Rotary embedding (rope-only); rotate first rope_dim dims + if rope_dim is None: + # Auto-compute: multiple of 6, capped at head_dim + self.rope_dim = max(6, (self.head_dim // 6) * 6) + self.rope_dim = min(self.rope_dim, self.head_dim) + else: + # Use custom rope_dim (must be multiple of 6 for 3D RoPE) + assert ( + rope_dim % 6 == 0 + ), f"rope_dim must be multiple of 6 for 3D RoPE, got {rope_dim}" + assert ( + rope_dim <= self.head_dim + ), f"rope_dim {rope_dim} exceeds head_dim {self.head_dim}" + self.rope_dim = rope_dim + + # ParticleRotaryEmbedding expects dim parameter but outputs dim//2 frequencies + # Since we need rope_dim output, we pass rope_dim*2 + self.rope3d = ParticleRotaryEmbedding( + dim=self.rope_dim * 2, + ) + # Relative position bias per head, from 3D delta + self.relpos_bias = nn.Linear(3, self.num_heads, bias=False) + + def _reshape_to_heads(self, x: torch.Tensor) -> torch.Tensor: + B, N, D = x.shape + x = x.view(B, N, self.num_heads, self.head_dim).permute( + 0, 2, 1, 3 + ) # (B,H,N,Hd) + return x + + def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: + B, H, N, Hd = x.shape + return x.permute(0, 2, 1, 3).contiguous().view(B, N, H * Hd) + + def forward( + self, + query: torch.Tensor, + key_value: torch.Tensor, + pos_q: torch.Tensor, + pos_k: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + B, Nq, _ = query.shape + _, Nk, _ = key_value.shape + + q = self._reshape_to_heads(self.q_proj(query)) # (B,H,Nq,Hd) + k = self._reshape_to_heads(self.k_proj(key_value)) # (B,H,Nk,Hd) + v = self._reshape_to_heads(self.v_proj(key_value)) # (B,H,Nk,Hd) + + # QK-normalization (RenderFormer): normalize Q and K before RoPE + if self.qk_norm: + q = self.q_norm(q) # (B,H,Nq,Hd) - applied per head + k = self.k_norm(k) # (B,H,Nk,Hd) - applied per head + + # RoPE: detect single-point (3D) vs multi-point (Nq*3 D) input + if pos_q is None or pos_k is None: + raise ValueError( + "CrossAttentionWithRoPE expects pos_q and pos_k (rope-only)" + ) + + # Compute RoPE embeddings based on input dimensionality + rope_q = self.rope3d(pos_q) # (B,Nq,rope_dim) + rope_k = self.rope3d(pos_k) # (B,Nk,rope_dim) + + # Expand to heads: (B,H,N,rope_dim) + rope_q = rope_q.unsqueeze(1).expand(B, self.num_heads, Nq, self.rope_dim) + rope_k = rope_k.unsqueeze(1).expand(B, self.num_heads, Nk, self.rope_dim) + + import pdb + + pdb.set_trace() # 1. check rot pass dim, check self.rope3d + # Apply rotation to q/k + q_rot = q[..., : self.rope_dim] + k_rot = k[..., : self.rope_dim] + q_pass = q[..., self.rope_dim :] + k_pass = k[..., self.rope_dim :] + q_rot = apply_rotary_emb(rope_q, q_rot) + k_rot = apply_rotary_emb(rope_k, k_rot) + q = torch.cat([q_rot, q_pass], dim=-1) + k = torch.cat([k_rot, k_pass], dim=-1) + + # Attention + attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale # (B,H,Nq,Nk) + + pos_q_3d = pos_q # (B, Nq, 3) + pos_k_3d = pos_k # (B, Nk, 3) + delta = pos_q_3d.unsqueeze(2) - pos_k_3d.unsqueeze(1) # (B, Nq, Nk, 3) + bias = self.relpos_bias(delta) # (B, Nq, Nk, H) + bias = bias.permute(0, 3, 1, 2).contiguous() # (B, H, Nq, Nk) + attn_scores = attn_scores + bias + if attn_mask is not None: + # attn_mask expected shape broadcastable to (B, 1, Nq, Nk) or (B,H,Nq,Nk) + attn_scores = attn_scores.masked_fill(attn_mask == 0, float("-inf")) + attn_probs = attn_scores.softmax(dim=-1) + attn_probs = self.dropout(attn_probs) + + out = torch.matmul(attn_probs, v) # (B,H,Nq,Hd) + out = self._merge_heads(out) # (B,Nq,D) + out = self.out_proj(out) + return out + + +from .merge import merge_source, merge_wavg, k_partite_soft_matching + + +class MultiHeadAttentionWithTOME(nn.Module): + + def __init__( + self, + query_dim, + num_heads, + ffn_hidden_dim, + kv_dim=None, + dropout=0.1, + bias=True, + bias_kv=False, + activation: str = "swiglu", + norm_type="layer_norm", + disable_q_norm=False, + disable_kv_norm=False, + qk_norm=False, + add_self_attn: bool = False, + ): + super().__init__() + self.apply_rope_cossin = apply_rotary_emb_cossin + + self.num_heads = num_heads + self.is_self_attn = kv_dim is None + kv_dim = query_dim if kv_dim is None else kv_dim + + self.in_proj = nn.Linear(query_dim, 3 * query_dim, bias=bias) + self.out_proj = nn.Linear(query_dim, query_dim, bias=bias) + + if qk_norm: + if norm_type == "layer_norm": + norm_module = nn.LayerNorm + elif norm_type == "rms_norm": + norm_module = RMSNorm + else: + raise ValueError( + "Unsupported normalization type. Choose from 'layer_norm' and 'rms_norm'." + ) + self.q_norm = norm_module(query_dim, eps=EPS) + self.k_norm = norm_module(query_dim, eps=EPS) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + assert bias_kv is False, "Bias for key and value is not supported for now" + self.query_norm = ( + norm_module(query_dim, eps=EPS) if not disable_q_norm else nn.Identity() + ) + + if activation == "swiglu": + self.ffn = FeedForwardSwiGLU( + query_dim, hidden_dim=ffn_hidden_dim, dropout=dropout, bias=bias + ) + elif activation == "gelu": + self.ffn = FeedForwardGeLU( + query_dim, hidden_dim=ffn_hidden_dim, dropout=dropout, bias=bias + ) + else: + raise ValueError( + "Unsupported activation function. Choose from 'gelu' and 'swiglu'." + ) + self.ffn_norm = norm_module(query_dim, eps=EPS) + + def forward( + self, + query, + num_to_remove_tokens, + size_token=None, # size of each token + src_key_padding_mask=None, + rope_cos=None, + rope_sin=None, + is_trace_source=False, + trace_source=None, + force_sdpa=False, + k_partite=-1, + ): + + # src_key_padding_mask: (B, N), key padding mask, things you want to attend to is True + bs, src_len = query.shape[0], query.shape[1] + + q = self.query_norm(query) + kv = q + ctx_len = kv.shape[1] + + # multihead self attention + q, k, v = self.in_proj(q).chunk(3, dim=-1) + + # qk normalization + q = self.q_norm(q).type(v.dtype) + k = self.k_norm(k).type(v.dtype) + + q = q.view(bs, src_len, self.num_heads, -1).transpose( + 1, 2 + ) # (bs, num_heads, src_len, head_dim) + k = k.view(bs, ctx_len, self.num_heads, -1).transpose(1, 2) + v = v.view(bs, ctx_len, self.num_heads, -1).transpose(1, 2) + + # apply rope + if rope_cos is not None: + q, k = self.apply_rope_cossin(q, k, rope_cos, rope_sin) + + if ATTN == "sdpa" or force_sdpa: + # create attention mask + if src_key_padding_mask is not None: + # things you want to attend to is True + assert src_key_padding_mask.shape == ( + bs, + ctx_len, + ), f"expecting key_padding_mask shape of {(bs, ctx_len)}, but got {src_key_padding_mask.shape}" + # but now, things you want to attend to is True + src_key_padding_mask = torch.zeros_like( + src_key_padding_mask, dtype=torch.bool + ).masked_fill_(src_key_padding_mask, True) + attn_mask = ( + src_key_padding_mask.view(bs, 1, 1, ctx_len) + .expand(-1, self.num_heads, -1, -1) + .reshape(bs * self.num_heads, 1, ctx_len) + ) + attn_mask = attn_mask.view(bs, self.num_heads, -1, ctx_len) + else: + attn_mask = None + + attn_output = ( + F.scaled_dot_product_attention( + query=q.type(v.dtype), + key=k.type(v.dtype), + value=v, + attn_mask=attn_mask, + ) + .transpose(1, 2) + .contiguous() + .view(bs, src_len, -1) + ) + elif ATTN == "flash_attn": + # self-attn + if src_key_padding_mask is not None: + q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, _ = unpad_input( + q.transpose(1, 2), src_key_padding_mask + ) + k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, _ = unpad_input( + k.transpose(1, 2), src_key_padding_mask + ) + v_unpad, indices_v, cu_seqlens_v, max_seqlen_v, _ = unpad_input( + v.transpose(1, 2), src_key_padding_mask + ) + qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1) + out_unpad = flash_attn_varlen_qkvpacked_func( + qkv_unpad, + cu_seqlens_q, + max_seqlen_q, + ) + attn_output = ( + pad_input(out_unpad, indices_q, bs, src_len) + .contiguous() + .view(bs, src_len, -1) + ) + else: + qkv = torch.stack( + [ + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + ], + dim=2, + ) + attn_output = ( + flash_attn_qkvpacked_func(qkv).contiguous().view(bs, src_len, -1) + ) + else: + raise ValueError( + "Unsupported attention type. Choose from 'flash_attn' and 'sdpa'." + ) + + x_attn = self.out_proj(attn_output) + metric = k.mean(dim=1) # (B, N, head_dim) + + # merge + # TODO the self attn does not include size_token yet + query = query + self.dropout(x_attn) + if num_to_remove_tokens > 0: + assert k_partite >= 2, "k_partite should be at least 2" + merge, _ = k_partite_soft_matching( + metric, + k=k_partite, + class_token=False, + distill_token=False, + ) + + if is_trace_source: + trace_source = merge_source(merge, query, trace_source) + + query, size_token = merge_wavg(merge, query, size_token) + + x = query + self.dropout(self.ffn(self.ffn_norm(query))) + return x, size_token, trace_source diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/layers/merge.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/layers/merge.py new file mode 100644 index 000000000..85511c07e --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/layers/merge.py @@ -0,0 +1,457 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- + +import math +from typing import Callable, Tuple, Union, List, Optional + +import torch + + +def do_nothing(x, mode=None): + return x + + +def bipartite_soft_matching( + metric: torch.Tensor, + r: int, + class_token: bool = False, + distill_token: bool = False, +) -> Tuple[Callable, Callable]: + """ + Applies ToMe with a balanced matching set (50%, 50%). + + Input size is [batch, tokens, channels]. + r indicates the number of tokens to remove (max 50% of tokens). + + Extra args: + - class_token: Whether or not there's a class token. + - distill_token: Whether or not there's also a distillation token. + + When enabled, the class token and distillation tokens won't get merged. + """ + protected = 0 + if class_token: + protected += 1 + if distill_token: + protected += 1 + + # We can only reduce by a maximum of 50% tokens + t = metric.shape[1] + r = min(r, (t - protected) // 2) + + if r <= 0: + return do_nothing, do_nothing + + with torch.no_grad(): + metric = metric / metric.norm(dim=-1, keepdim=True) + a, b = metric[..., ::2, :], metric[..., 1::2, :] + scores = a @ b.transpose(-1, -2) + + if class_token: + scores[..., 0, :] = -math.inf + if distill_token: + scores[..., :, 0] = -math.inf + + node_max, node_idx = scores.max(dim=-1) + edge_idx = node_max.argsort(dim=-1, descending=True)[..., None] + + unm_idx = edge_idx[..., r:, :] # Unmerged Tokens + src_idx = edge_idx[..., :r, :] # Merged Tokens + dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx) + + if class_token: + # Sort to ensure the class token is at the start + unm_idx = unm_idx.sort(dim=1)[0] + + def merge(x: torch.Tensor, mode="mean") -> torch.Tensor: + src, dst = x[..., ::2, :], x[..., 1::2, :] + n, t1, c = src.shape + unm = src.gather(dim=-2, index=unm_idx.expand(n, t1 - r, c)) + src = src.gather(dim=-2, index=src_idx.expand(n, r, c)) + dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode) + + if distill_token: + return torch.cat([unm[:, :1], dst[:, :1], unm[:, 1:], dst[:, 1:]], dim=1) + else: + return torch.cat([unm, dst], dim=1) + + def unmerge(x: torch.Tensor) -> torch.Tensor: + unm_len = unm_idx.shape[1] + unm, dst = x[..., :unm_len, :], x[..., unm_len:, :] + n, _, c = unm.shape + + src = dst.gather(dim=-2, index=dst_idx.expand(n, r, c)) + + out = torch.zeros(n, metric.shape[1], c, device=x.device, dtype=x.dtype) + + out[..., 1::2, :] = dst + out.scatter_(dim=-2, index=(2 * unm_idx).expand(n, unm_len, c), src=unm) + out.scatter_(dim=-2, index=(2 * src_idx).expand(n, r, c), src=src) + + return out + + return merge, unmerge + + +def k_partite_soft_matching( + metric: torch.Tensor, + k: int, + r: Optional[int] = None, + class_token: bool = False, + distill_token: bool = False, +) -> Tuple[Callable, Callable]: + """ + ToMe-style bipartite soft matching generalized to k partitions. + + Semantics (matches `bipartite_soft_matching`, generalized): + - Build two sets: A (sources) and B (sinks). + * We choose B = { indices i : i % k == b_mod } and A = complement. + * By default b_mod = k-1, so B are the "k-1, 2k-1, ..." tokens. + * If distill_token=True, we set b_mod so that the distill token (index 1) + lies in B and is placed as the FIRST column of B (never a merge target). + * If class_token=True, the class token (index 0) is in A and is never merged. + - Score with cosine: scores = A @ B^T + - For each A row, pick its best B column; rank A by that score. + - Merge the top-r A rows into their chosen B rows via scatter_reduce. + - Output layout mirrors ToMe: + * if distill_token: [unm[:1] (class), dst[:1] (distill), unm[1:], dst[1:]] + * else: [unm, dst] + where `unm` are unmerged A tokens (class first if present) and `dst` are the + updated B tokens after merging. + + Args: + metric: [B, T, C] features used only for computing the selection (indices). + k: number of partitions (k >= 2). + r: number of A tokens to merge (if None -> merge as many A as possible + while honoring protected tokens; i.e., all A except class if protected). + class_token, distill_token: same meaning and handling as in bipartite_soft_matching. + + Returns: + merge(x): [B, T, C] -> [B, (|A|-r) + |B|, C] + unmerge(y): [B, (|A|-r) + |B|, C] -> [B, T, C] + """ + assert k >= 2, "k must be >= 2" + Bsz, T, C = metric.shape + + # Decide which modulo class forms B so distill ∈ B and class ∉ B + if distill_token: + b_mod = 1 % k # make sure distill (index 1) is a sink + else: + b_mod = (k - 1) % k + if class_token and b_mod == 0: + # ensure class (index 0) stays in A + b_mod = (b_mod + 1) % k if k > 2 else 1 + + # Build A/B index lists (ascending, deterministic) + all_idx = torch.arange(T, device=metric.device, dtype=torch.long) + is_B = (all_idx % k) == b_mod + B_idx = all_idx[is_B] + A_idx = all_idx[~is_B] + TA, TB = A_idx.numel(), B_idx.numel() + + # Clamp r like ToMe, but w.r.t. size(A) and protected tokens + protected_in_A = 1 if (class_token and (TA > 0) and (A_idx[0].item() == 0)) else 0 + r_max = max(0, TA - protected_in_A) + if r is None: + r_use = r_max + else: + r_use = max(0, min(int(r), r_max)) + + if r_use <= 0: + # Nothing to merge → return identity closures + def _id(x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + return x + + return _id, _id + + with torch.no_grad(): + # Cosine-normalize for scoring (same as ToMe) + metric_n = metric / (metric.norm(dim=-1, keepdim=True) + 1e-12) + + # Gather A/B feature matrices + a = metric_n.index_select(dim=1, index=A_idx) # [B, TA, C] + b = metric_n.index_select(dim=1, index=B_idx) # [B, TB, C] + + # Reorder B so distill column is first (like ToMe puts distill at column 0 of B) + if distill_token: + where_d = (B_idx == 1).nonzero(as_tuple=False).flatten() + if where_d.numel() > 0: + dpos = int(where_d.item()) + perm = torch.cat( + [ + torch.tensor([dpos], device=b.device), + torch.arange(TB, device=b.device)[ + torch.arange(TB, device=b.device) != dpos + ], + ] + ) + b = b.index_select(1, perm) + B_idx = B_idx.index_select(0, perm) + + # Scores: A x B (cosine) + scores = a @ b.transpose(-1, -2) # [B, TA, TB] + + # Protect class (row in A) and distill (column in B) + if class_token: + # Because A_idx is ascending, if class is in A it must be at row 0. + if TA > 0 and A_idx[0].item() == 0: + scores[..., 0, :] = -math.inf + if distill_token: + if TB > 0 and B_idx[0].item() == 1: + scores[..., :, 0] = -math.inf + + # Best B per A, then rank A by its best score + node_max, node_idx = scores.max(dim=-1) # [B, TA], [B, TA] + edge_idx = node_max.argsort(dim=-1, descending=True)[..., None] # [B, TA, 1] + + unm_idx = edge_idx[..., r_use:, :] # A-rows we keep as individual tokens + src_idx = edge_idx[..., :r_use, :] # A-rows we will merge into B + dst_idx = node_idx[..., None].gather( + dim=-2, index=src_idx + ) # [B, r_use, 1] (B-row per src) + + if class_token and (TA > 0) and (A_idx[0].item() == 0): + # Keep the class token first in the unmerged block + unm_idx = unm_idx.sort(dim=1)[0] + + # ---- Closures (autograd preserved for x) ---- + + def merge(x: torch.Tensor, mode: str = "mean") -> torch.Tensor: + """ + x: [B, T, C] -> y: [B, (TA - r) + TB, C] + """ + src_all = x.index_select(1, A_idx) # [B, TA, C] + dst_all = x.index_select(1, B_idx) # [B, TB, C] + + B, _, Cx = src_all.shape + + # Unmerged A tokens + unm = src_all.gather( + dim=-2, index=unm_idx.expand(B, TA - r_use, Cx) + ) # [B, TA-r, C] + + # Selected sources and their destination rows in B + src_sel = src_all.gather( + dim=-2, index=src_idx.expand(B, r_use, Cx) + ) # [B, r, C] + dst_new = dst_all.scatter_reduce( + -2, dst_idx.expand(B, r_use, Cx), src_sel, reduce=mode + ) + + if distill_token: + # [class], [distill], [other unmerged A], [other B sinks] + return torch.cat( + [unm[:, :1], dst_new[:, :1], unm[:, 1:], dst_new[:, 1:]], dim=1 + ) + else: + return torch.cat([unm, dst_new], dim=1) + + def unmerge(y: torch.Tensor) -> torch.Tensor: + """ + y: [B, (TA - r) + TB, C] -> x_hat: [B, T, C] + """ + B, Ylen, Cx = y.shape + unm_len = unm_idx.shape[1] + unm, dst_cur = ( + y[..., :unm_len, :], + y[..., unm_len:, :], + ) # split exactly like merge's concat + + # Reconstruct the merged A sources by pulling from their chosen B rows + src_rec = dst_cur.gather( + dim=-2, index=dst_idx.expand(B, r_use, Cx) + ) # [B, r, C] + + # Prepare a zero output and scatter everything back to original slots + out = torch.zeros(B, T, Cx, device=y.device, dtype=y.dtype) + + # Global indices for A (unmerged and merged) and B + A_idx_batched = A_idx.view(1, TA, 1).expand(B, TA, 1) # [B, TA, 1] + unm_pos = A_idx_batched.gather(dim=1, index=unm_idx) # [B, TA-r, 1] + src_pos = A_idx_batched.gather(dim=1, index=src_idx) # [B, r, 1] + B_idx_batched = B_idx.view(1, TB, 1).expand(B, TB, 1) # [B, TB, 1] + + # Scatter B block (dst_cur) to B positions + out.scatter_(dim=1, index=B_idx_batched.expand(B, TB, Cx), src=dst_cur) + # Scatter unmerged A + out.scatter_(dim=1, index=unm_pos.expand(B, TA - r_use, Cx), src=unm) + # Scatter reconstructed merged A + out.scatter_(dim=1, index=src_pos.expand(B, r_use, Cx), src=src_rec) + + return out + + return merge, unmerge + + +def kth_bipartite_soft_matching( + metric: torch.Tensor, k: int +) -> Tuple[Callable, Callable]: + """ + Applies ToMe with the two sets as (every kth element, the rest). + If n is the number of tokens, resulting number of tokens will be n // z. + + Input size is [batch, tokens, channels]. + z indicates the stride for the first set. + z = 2 is equivalent to regular bipartite_soft_matching with r = 0.5 * N + """ + if k <= 1: + return do_nothing, do_nothing + + def split(x): + t_rnd = (x.shape[1] // k) * k + x = x[:, :t_rnd, :].view(x.shape[0], -1, k, x.shape[2]) + a, b = ( + x[:, :, : (k - 1), :].contiguous().view(x.shape[0], -1, x.shape[-1]), + x[:, :, (k - 1), :], + ) + return a, b + + with torch.no_grad(): + metric = metric / metric.norm(dim=-1, keepdim=True) + a, b = split(metric) + r = a.shape[1] + scores = a @ b.transpose(-1, -2) + + _, dst_idx = scores.max(dim=-1) + dst_idx = dst_idx[..., None] + + def merge(x: torch.Tensor, mode="mean") -> torch.Tensor: + src, dst = split(x) + n, _, c = src.shape + dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode) + + return dst + + def unmerge(x: torch.Tensor) -> torch.Tensor: + n, _, c = x.shape + dst = x + + src = dst.gather(dim=-2, index=dst_idx.expand(n, r, c)).to(x.dtype) + + src = src.view(n, -1, (k - 1), c) + dst = dst.view(n, -1, 1, c) + + out = torch.cat([src, dst], dim=-2) + out = out.contiguous().view(n, -1, c) + + return out + + return merge, unmerge + + +def random_bipartite_soft_matching( + metric: torch.Tensor, r: int +) -> Tuple[Callable, Callable]: + """ + Applies ToMe with the two sets as (r chosen randomly, the rest). + Input size is [batch, tokens, channels]. + + This will reduce the number of tokens by r. + """ + if r <= 0: + return do_nothing, do_nothing + + with torch.no_grad(): + B, N, _ = metric.shape + rand_idx = torch.rand(B, N, 1, device=metric.device).argsort(dim=1) + + a_idx = rand_idx[:, :r, :] + b_idx = rand_idx[:, r:, :] + + def split(x): + C = x.shape[-1] + a = x.gather(dim=1, index=a_idx.expand(B, r, C)) + b = x.gather(dim=1, index=b_idx.expand(B, N - r, C)) + return a, b + + metric = metric / metric.norm(dim=-1, keepdim=True) + a, b = split(metric) + scores = a @ b.transpose(-1, -2) + + _, dst_idx = scores.max(dim=-1) + dst_idx = dst_idx[..., None] + + def merge(x: torch.Tensor, mode="mean") -> torch.Tensor: + src, dst = split(x) + C = src.shape[-1] + dst = dst.scatter_reduce(-2, dst_idx.expand(B, r, C), src, reduce=mode) + + return dst + + def unmerge(x: torch.Tensor) -> torch.Tensor: + C = x.shape[-1] + dst = x + src = dst.gather(dim=-2, index=dst_idx.expand(B, r, C)) + + out = torch.zeros(B, N, C, device=x.device, dtype=x.dtype) + + out.scatter_(dim=-2, index=a_idx.expand(B, r, C), src=src) + out.scatter_(dim=-2, index=b_idx.expand(B, N - r, C), src=dst) + + return out + + return merge, unmerge + + +def merge_wavg( + merge: Callable, x: torch.Tensor, size: torch.Tensor = None +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Applies the merge function by taking a weighted average based on token size. + Returns the merged tensor and the new token sizes. + """ + if size is None: + size = torch.ones_like(x[..., 0, None]) + + x = merge(x * size, mode="sum") + size = merge(size, mode="sum") + + x = x / size + return x, size + + +def merge_source( + merge: Callable, x: torch.Tensor, source: torch.Tensor = None +) -> torch.Tensor: + """ + For source tracking. Source is an adjacency matrix between the initial tokens and final merged groups. + x is used to find out how many tokens there are in case the source is None. + """ + if source is None: + n, t, _ = x.shape + source = torch.eye(t, device=x.device)[None, ...].expand(n, t, t) + + source = merge(source, mode="amax") + return source + + +def parse_r(num_layers: int, r: Union[List[int], Tuple[int, float], int]) -> List[int]: + """ + Process a constant r or r schedule into a list for use internally. + + r can take the following forms: + - int: A constant number of tokens per layer. + - Tuple[int, float]: A pair of r, inflection. + Inflection describes there the the reduction / layer should trend + upward (+1), downward (-1), or stay constant (0). A value of (r, 0) + is as providing a constant r. (r, -1) is what we describe in the paper + as "decreasing schedule". Any value between -1 and +1 is accepted. + - List[int]: A specific number of tokens per layer. For extreme granularity. + """ + inflect = 0 + if isinstance(r, list): + if len(r) < num_layers: + r = r + [0] * (num_layers - len(r)) + return list(r) + elif isinstance(r, tuple): + r, inflect = r + + min_val = int(r * (1.0 - inflect)) + max_val = 2 * r - min_val + step = (max_val - min_val) / (num_layers - 1) + + return [int(min_val + step * i) for i in range(num_layers)] diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/particle_network_cross_attn_feat.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/particle_network_cross_attn_feat.py new file mode 100644 index 000000000..19ef37f8f --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/particle_network_cross_attn_feat.py @@ -0,0 +1,548 @@ +import torch +import torch.nn as nn +import numpy as np +import time +from .cross_attn_feat import CrossAttentionWeightedPoswithFeature +from .encodings.rope import ParticleRotaryEmbedding, freqs_to_cos_sin +from .super_particle_layers import LocalFeatureExtractor, SuperParticleDecoder +from ._data import uses_obstacle_features, uses_topology_neighbors + + +class ParticleNetworkCrossAttnLocalFeature(torch.nn.Module): + """ + Particle Network with Slot Attention for adaptive super-particle generation. + + Architecture: + 1. Initial feature extraction via continuous convolutions + 2. Slot attention to generate super-particles (adaptive sampling) + 3. Super-particle position prediction + 4. Self-attention among super-particles + 5. Cross-attention from all particles to super-particles with RoPE + 6. Position correction prediction + """ + + def __init__( + self, + # ============= Continuous Convolution Parameters ============= + kernel_size=[4, 4, 4], + cconv_embedding_dim=32, + radius_scale=1.5, + data=None, + use_topology_conv=None, + topology_search_radius=None, + topology_blend_alpha=None, + coordinate_mapping="ball_to_cube_volume_preserving", + interpolation="linear", + use_window=True, + particle_radius=0.025, + # ============= Position Encoding Parameters ============= + particle_position_rope_dim=96, # Must be divisible by 6 for 3D positions + # ============= Slot Attention Parameters ============= + slot_compression_ratio=10, # N_particles / N_slots ratio + slot_attn_iters=3, + slot_attn_heads=6, + # ============= Decoder Parameters ============= + num_decoder_layers=3, + decoder_attn_heads=4, + decoder_attn_ffn_hidden_dim=384, + decoder_attn_dropout=0.1, + # ============= Output Network Parameters ============= + output_hidden_dim=128, + output_layers=2, + # ============= Physics Parameters ============= + timestep=1 / 50, + gravity=(0, -9.81, 0), + other_feats_channels=0, + obstacle_feats_channels=3, + # ============= Debug Parameters ============= + verbose=False, + velhead=False, + ): + super().__init__() + + # ============= Store Basic Configuration ============= + self.velhead = velhead + self.data = data + self.topology_search_radius = topology_search_radius + self.verbose = verbose + self.timestep = timestep + self.gravity_values = list(gravity) + self.slot_compression_ratio = slot_compression_ratio + self.decoder_enable_self_attn = self.data not in { + "cloth_no_self_attention", + "cloth_no_all", + } + self.decoder_enable_ffn = self.data not in {"cloth_no_ffn", "cloth_no_all"} + + if use_topology_conv is None: + self.use_topology_conv = uses_topology_neighbors(self.data) + else: + self.use_topology_conv = bool(use_topology_conv) + + # ============ Store RoPE Configuration ============= + self.particle_position_rope_dim = particle_position_rope_dim + + assert ( + self.particle_position_rope_dim % 6 == 0 + ), "particle_position_rope_dim must be divisible by 6 for 3D positions" + + assert ( + self.particle_position_rope_dim + <= cconv_embedding_dim * 3 // decoder_attn_heads + ), "RoPE dimension per head must be smaller than head dimension" + + # ============= Continuous Convolution Configuration ============= + self.local_feature_extractor = LocalFeatureExtractor( + kernel_size=kernel_size, + cconv_embedding_dim=cconv_embedding_dim, + data=data, + use_topology_conv=self.use_topology_conv, + topology_search_radius=topology_search_radius, + topology_blend_alpha=topology_blend_alpha, + radius_scale=radius_scale, + coordinate_mapping=coordinate_mapping, + interpolation=interpolation, + use_window=use_window, + particle_radius=particle_radius, + other_feats_channels=other_feats_channels, + obstacle_feats_channels=obstacle_feats_channels, + ) + + # ============= Calculate Feature Dimensions ============= + # Combined dimension: obstacle conv + fluid conv + dense layer + if self.data == "local_fluid": + combined_feat_dim = cconv_embedding_dim * 3 + elif uses_obstacle_features(self.data): + combined_feat_dim = cconv_embedding_dim * 3 + else: + combined_feat_dim = cconv_embedding_dim * 2 + + # ============= Initialize Slot Attention ============= + self.super_particle_cross_attn_feat = CrossAttentionWeightedPoswithFeature( + dim=combined_feat_dim, + filter_extent=self.local_feature_extractor.filter_extent, + iters=slot_attn_iters, + n_heads=slot_attn_heads, + eps=1e-8, + ) + + # ============= Initialize RoPE Encodings ============= + # RoPE for particle positions + self.particle_rope = ParticleRotaryEmbedding( + dim=self.particle_position_rope_dim + ) + + # RoPE for super-particle positions + self.super_particle_rope = ParticleRotaryEmbedding( + dim=self.particle_position_rope_dim + ) + + # ============= Initialize Super-Particle Decoder ============= + self.super_particle_decoder = SuperParticleDecoder( + num_layers=num_decoder_layers, + query_particle_feat_dim=combined_feat_dim, + super_particle_feat_dim=combined_feat_dim, + num_heads=decoder_attn_heads, + ffn_hidden_dim=decoder_attn_ffn_hidden_dim, + dropout=decoder_attn_dropout, + bias=True, + activation="swiglu", + norm_type="rms_norm", + qk_norm=True, + enable_self_attn=self.decoder_enable_self_attn, + enable_ffn=self.decoder_enable_ffn, + ) + + # ============= Initialize Output Network ============= + self.output_network = self._build_output_network( + input_dim=combined_feat_dim, + hidden_dim=output_hidden_dim, + num_layers=output_layers, + ) + + # Output scaling factor for numerical stability + self.output_scale = 1.0 / 128 + + def apply_radius_override(self, radius_scale=None, particle_radius=None): + """ + Re-apply radius-dependent values after checkpoint loading. + This is needed because state_dict restores buffers such as + `local_feature_extractor.filter_extent_tensor`. + """ + lfe = self.local_feature_extractor + if radius_scale is not None: + lfe.radius_scale = float(radius_scale) + if particle_radius is not None: + lfe.particle_radius = float(particle_radius) + + new_filter_extent = np.float32(lfe.radius_scale * 6 * lfe.particle_radius) + lfe.filter_extent = new_filter_extent + + if hasattr(lfe, "filter_extent_tensor"): + lfe.filter_extent_tensor.data.fill_(float(new_filter_extent)) + + if hasattr(lfe, "topology_filter_extent_tensor"): + if getattr(lfe, "topology_search_radius", None) is None: + lfe.topology_filter_extent = new_filter_extent + else: + lfe.topology_filter_extent = np.float32( + lfe.radius_scale * 6 * float(lfe.topology_search_radius) + ) + lfe.topology_filter_extent_tensor.data.fill_( + float(lfe.topology_filter_extent) + ) + + if hasattr(self, "super_particle_cross_attn_feat"): + self.super_particle_cross_attn_feat.filter_extent = float(new_filter_extent) + + def _build_output_network(self, input_dim, hidden_dim, num_layers): + """Build the output MLP network for position correction.""" + layers = [] + current_dim = input_dim + + for i in range(num_layers): + if i < num_layers - 1: + # Hidden layers with activation and normalization + layers.extend( + [ + nn.Linear(current_dim, hidden_dim), + nn.ReLU(), + nn.LayerNorm(hidden_dim), + ] + ) + current_dim = hidden_dim + else: + # Final layer outputs 3D position corrections + if self.velhead: + layers.append(nn.Linear(current_dim, 6)) + else: + layers.append(nn.Linear(current_dim, 3)) + + return nn.Sequential(*layers) + + def integrate_pos_vel(self, pos1, vel1, external_force=None): + """ + Apply gravity and integrate position and velocity using Verlet integration. + + Uses semi-implicit Euler method for stable integration: + 1. Update velocity with gravity acceleration + 2. Update position using average of old and new velocity + + Args: + pos1: Current positions [N, 3] + vel1: Current velocities [N, 3] + + Returns: + pos2: Predicted positions after timestep [N, 3] + vel2: Predicted velocities after timestep [N, 3] + """ + dt = self.timestep + + if external_force is not None: + gravity = external_force.to(dtype=vel1.dtype, device=vel1.device) + else: + # Create gravity tensor fresh each time to avoid inplace operation issues + gravity = torch.tensor( + self.gravity_values, dtype=vel1.dtype, device=vel1.device + ) + # Broadcast gravity to match particle count + if vel1.dim() == 2: + gravity = gravity.unsqueeze(0).expand(vel1.size(0), -1) + + # Semi-implicit Euler integration + vel2 = vel1 + (dt * gravity) # Update velocity with gravity + pos2 = pos1 + (dt * (vel2 + vel1) / 2) # Update position with average velocity + return pos2, vel2 + + def compute_new_pos_vel(self, pos1, vel1, pos2, vel2, pos_correction): + """ + Apply position correction and compute corrected velocities. + + The correction is applied to the predicted position, and velocity + is recomputed based on the total displacement from the original position. + + Args: + pos1: Original positions at start of timestep [N, 3] + vel1: Original velocities at start of timestep [N, 3] + pos2: Predicted positions after integration [N, 3] + vel2: Predicted velocities after integration [N, 3] + pos_correction: Learned position corrections [N, 3] + + Returns: + pos: Corrected final positions [N, 3] + vel: Corrected final velocities [N, 3] + """ + dt = self.timestep + pos = pos2 + pos_correction # Apply correction to predicted position + vel = (pos - pos1) / dt # Recompute velocity from total displacement + return pos, vel + + def compute_correction( + self, + pos, + vel, + other_feats, + box=None, + box_feats=None, + k_partite: int = -1, + fixed_radius_search_hash_table=None, + batch_info=None, + topology_neighbors_index=None, + topology_neighbors_row_splits=None, + ): + """ + Compute position correction using slot attention and super-particle decoder. + + This is the main processing pipeline that: + 1. Finds particle neighbors + 2. Extracts local features via continuous convolutions + 3. Creates super-particles via slot attention + 4. Processes super-particles through the decoder + 5. Outputs position corrections + + Args: + pos: Particle positions [N, 3] + vel: Particle velocities [N, 3] + other_feats: Additional particle features [N, C] or None + box: Obstacle/boundary positions [M, 3] + box_feats: Obstacle/boundary features [M, Cb] + fixed_radius_search_hash_table: Pre-computed spatial hash table (optional) + batch_info: Optional batching information (currently unused) + + Returns: + Position correction tensor [N, 3] + """ + t2 = time.time() + + particle_features, num_fluid_neighbors, _ = self.local_feature_extractor( + inp_pos=pos, + query_pos=pos, + inp_vel=vel, + other_feats=other_feats, + box=box, + box_feats=box_feats, + fixed_radius_search_hash_table=fixed_radius_search_hash_table, + batch_info=batch_info, + topology_neighbors_index=topology_neighbors_index, + topology_neighbors_row_splits=topology_neighbors_row_splits, + ) + + if num_fluid_neighbors is None: + self.num_fluid_neighbors = torch.zeros( + particle_features.shape[0], + dtype=particle_features.dtype, + device=particle_features.device, + ) + else: + self.num_fluid_neighbors = num_fluid_neighbors # for loss calculation + + t3 = time.time() + if self.verbose: + print(f"Feature extraction time: {t3 - t2:.4f} seconds") + import pdb + + pdb.set_trace() + + # ============= Step 3: Slot Attention for Super-Particle Generation ============= + # Convert to batch format for slot attention + if batch_info is None: + particle_features_batch = particle_features.unsqueeze(0) # [1, N, D] + pos_batch = pos.unsqueeze(0) # [1, N, 3] + else: + Batch = batch_info["batch_size"] + particle_features_batch = particle_features.reshape( + Batch, -1, particle_features.shape[-1] + ) # [B, N/B, D] + pos_batch = pos.reshape(Batch, -1, 3) # [B, N/B, 3] + + # Generate super-particles through iterative slot attention + super_particle_features_batch, super_particle_positions_batch = ( + self.super_particle_cross_attn_feat( + inputs_feat_batch=particle_features_batch, + input_pos_batch=pos_batch, + k_partite=k_partite, + ) + ) # Output: [1, num_slots, D] + + t4 = time.time() + if self.verbose: + print(f"Slot attention time: {t4 - t3:.4f} seconds") + import pdb + + pdb.set_trace() + + # ============= Step 4: Compute RoPE Embeddings ============= + + # Generate RoPE for particle positions + rope_freqs = self.particle_rope(pos_batch) + particle_rope_cos, particle_rope_sin = freqs_to_cos_sin( + rope_freqs, head_dim=self.super_particle_decoder.head_dim + ) + + # Generate RoPE for super-particle positions + rope_super_particle_freqs = self.super_particle_rope( + super_particle_positions_batch + ) + super_particle_rope_cos, super_particle_rope_sin = freqs_to_cos_sin( + rope_super_particle_freqs, head_dim=self.super_particle_decoder.head_dim + ) + + t5 = time.time() + if self.verbose: + print(f"RoPE computation time: {t5 - t4:.4f} seconds") + import pdb + + pdb.set_trace() + + # ============= Step 6: Super-Particle Decoder Processing ============= + # Process through decoder with self-attention and cross-attention layers + output_features = self.super_particle_decoder( + particles_feat=particle_features_batch, # query + super_particles_feat=super_particle_features_batch, # kv + rope_cos=particle_rope_cos, + rope_sin=particle_rope_sin, + rope_ctx_cos=super_particle_rope_cos, + rope_ctx_sin=super_particle_rope_sin, + ) + + t6 = time.time() + if self.verbose: + print(f"Decoder processing time: {t6 - t5:.4f} seconds") + import pdb + + pdb.set_trace() + + # ============= Step 7: Generate Position Corrections ============= + # Pass decoder output through final network to get position corrections + out = self.output_network(output_features) + + if batch_info is None: + out = out.squeeze(0) + + if self.velhead: + pos_correction = out[:, :3] + vel_correction = out[:, 3:] + pos_correction = pos_correction * self.output_scale + vel_correction = vel_correction * (self.output_scale) + else: + pos_correction = out + pos_correction = pos_correction * self.output_scale + else: + if self.velhead: + pos_correction = out[:, :, :3].reshape(-1, 3) + vel_correction = out[:, :, 3:].reshape(-1, 3) + pos_correction = pos_correction * self.output_scale + vel_correction = vel_correction * (self.output_scale) + else: + pos_correction = out.reshape(-1, 3) + pos_correction = pos_correction * self.output_scale + + if self.velhead: + return pos_correction, vel_correction, super_particle_positions_batch + else: + return pos_correction, super_particle_positions_batch + + def forward( + self, + inputs, + fixed_radius_search_hash_table=None, + batch_info=None, + return_super_particle_positions=False, + k_partite: int = -1, + topology_neighbors_index=None, + topology_neighbors_row_splits=None, + external_force=None, + ): + """ + Compute one simulation timestep. + + This is the main entry point that orchestrates the full simulation step: + 1. Integrate positions/velocities with physics (gravity) + 2. Compute learned corrections via neural network + 3. Apply corrections to get final positions/velocities + + Args: + inputs: Tuple of (pos, vel, feats, box, box_feats) where: + - pos: Particle positions [N, 3] + - vel: Particle velocities [N, 3] + - feats: Additional particle features [N, C] or None + - box: Obstacle/boundary positions [M, 3] + - box_feats: Obstacle/boundary features [M, Cb] + fixed_radius_search_hash_table: Pre-computed spatial hash (optional) + batch_info: Batching information (currently unused) + return_super_particle_positions: Whether to return super-particle info + compute_jsd_loss: Whether to compute JSD loss (saves computation if False) + + Returns: + Tuple of (pos_corrected, vel_corrected): + - pos_corrected: Corrected positions after timestep [N, 3] + - vel_corrected: Corrected velocities after timestep [N, 3] + """ + # Unpack input tensors + if len(inputs) == 5: + pos, vel, feats, box, box_feats = inputs + elif len(inputs) == 7: + ( + pos, + vel, + feats, + box, + box_feats, + topology_neighbors_index_inp, + topology_neighbors_row_splits_inp, + ) = inputs + if topology_neighbors_index is None: + topology_neighbors_index = topology_neighbors_index_inp + if topology_neighbors_row_splits is None: + topology_neighbors_row_splits = topology_neighbors_row_splits_inp + else: + raise ValueError( + f"Unexpected input tuple length: {len(inputs)}. Expected 5 or 7." + ) + + # Step 1: Physics-based integration (gravity + velocity) + if self.velhead: + pos2, vel2 = self.integrate_pos_vel(pos, vel, external_force=external_force) + else: + pos2, vel2 = pos, vel + + # Step 2: Compute learned corrections using neural network + if self.velhead: + pos_correction, vel_correction, super_particle_positions_batch = ( + self.compute_correction( + pos2, + vel2, + feats, + box, + box_feats, + k_partite=k_partite, + fixed_radius_search_hash_table=fixed_radius_search_hash_table, + batch_info=batch_info, + topology_neighbors_index=topology_neighbors_index, + topology_neighbors_row_splits=topology_neighbors_row_splits, + ) + ) + + pos2_corrected = pos2 + pos_correction + vel2_corrected = vel2 + vel_correction + else: + pos_correction, super_particle_positions_batch = self.compute_correction( + pos2, + vel2, + feats, + box, + box_feats, + k_partite=k_partite, + fixed_radius_search_hash_table=fixed_radius_search_hash_table, + batch_info=batch_info, + topology_neighbors_index=topology_neighbors_index, + topology_neighbors_row_splits=topology_neighbors_row_splits, + ) + pos2_corrected, vel2_corrected = self.compute_new_pos_vel( + pos, vel, pos2, vel2, pos_correction + ) + + if return_super_particle_positions: + return pos2_corrected, vel2_corrected, super_particle_positions_batch + + return pos2_corrected, vel2_corrected diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/super_particle_layers.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/super_particle_layers.py new file mode 100644 index 000000000..b2b940d15 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_network/super_particle_layers.py @@ -0,0 +1,1004 @@ +import torch +import torch.nn as nn +from .layers.attention import MultiHeadAttention, FeedForwardGeLU, FeedForwardSwiGLU +from .encodings.rms_norm import RMSNorm +from .continuous_conv import ContinuousConv, ParticleRadiusResearch, FixedRadiusSearch +from open3d.ml.torch.ops import reduce_subarrays_sum +import numpy as np +import math +from ._data import uses_obstacle_features + + +class CubicSplineKernel: + """ + Cubic Spline Kernel implementation based on SPlisHSPlasH + W(r,h) = (8/(pi*h^3)) * { + 6*q^3 - 6*q^2 + 1 if 0 <= q <= 0.5 + 2*(1-q)^3 if 0.5 < q <= 1 + 0 if q > 1 + } + where q = ||r|| / h + """ + + def __init__(self, particle_radius=0.025): + h = 2.0 * particle_radius # Kernel support radius + self.h = h + self.h2 = h * h + self.h3 = h * h * h + + # Normalization constants + self.k = 8.0 / (math.pi * self.h3) + self.l = 48.0 / (math.pi * self.h3) + + +class SPHVelocityInterpolator: + """ + SPH Velocity Interpolator using the formulation from SPlisHSPlasH + + The velocity at query point x_i is interpolated as: + v_i = sum_j (m_j / rho_j) * v_j * W(x_i - x_j, h) + + For simplicity, we assume uniform mass and density (m_j/rho_j = constant) + """ + + def __init__(self, kernel_radius): + """ + Initialize the interpolator + + Args: + particle_positions: numpy array of shape (N, 3) - particle positions + particle_velocities: numpy array of shape (N, 3) - particle velocities + kernel_radius: float - SPH kernel support radius + """ + # TODO: check if this is correct + self.kernel_radius = kernel_radius + self.kernel = CubicSplineKernel(kernel_radius) + self.fixed_radius_search = FixedRadiusSearch( + metric="L2", ignore_query_point=False, return_distances=True + ) + + def interpolate_velocity( + self, query_points, input_positions, input_velocities, input_pos_hash_table=None + ): + """ + Args: + query_points: (Q, 3) torch.float + input_positions: (N, 3) torch.float + input_velocities: (N, 3) torch.float + Returns: + (Q, 3) torch.float interpolated velocities + """ + # Neighbor search (assumes torch tensors returned) + if input_pos_hash_table is not None: + nns = self.fixed_radius_search( + input_positions, + queries=query_points, + radius=self.kernel_radius, + hash_table=input_pos_hash_table, + ) + else: + nns, input_pos_hash_table = self.fixed_radius_search( + input_positions, + queries=query_points, + radius=self.kernel_radius, + return_hash_table=True, + ) + + # Flattened neighbor info + d = nns.neighbors_distance.sqrt() # (M,) distances (because input were squared) + nbr_idx = nns.neighbors_index # (M,) long + row_splits = nns.neighbors_row_splits # (Q+1,) long + + device = input_velocities.device + dtype = input_velocities.dtype + + # Kernel constants + h = torch.as_tensor(self.kernel.h, device=device, dtype=dtype) + k = torch.as_tensor(self.kernel.k, device=device, dtype=dtype) + + # Vectorized cubic spline weights W(r,h) + q = d.to(dtype) / h # (M,) + w = torch.zeros_like(q) + + m1 = q <= 0.5 + m2 = (q > 0.5) & (q <= 1.0) + + q2 = q * q + q3 = q2 * q + w = w.masked_scatter(m1, k * (6.0 * q3[m1] - 6.0 * q2[m1] + 1.0)) + w = w.masked_scatter(m2, k * 2.0 * (1.0 - q[m2]) ** 3) + # values where q > 1 remain 0 + + # Neighbor velocity contributions (M, 3) + contrib = input_velocities[nbr_idx] * w.unsqueeze(-1) + + # Build query index per neighbor via row_splits + lengths = (row_splits[1:] - row_splits[:-1]).to(torch.long) # (Q,) + Q = lengths.numel() + q_idx_per_nbr = torch.repeat_interleave( + torch.arange(Q, device=device, dtype=torch.long), lengths + ) # (M,) + + # Accumulate per-query sums: out[q] += contrib[m] where q = q_idx_per_nbr[m] + out = torch.zeros((Q, 3), device=device, dtype=dtype) + out.scatter_add_(0, q_idx_per_nbr.view(-1, 1).expand(-1, 3), contrib) + + # Optional: Shepard-style normalization for smoother interpolation + # sum_w[q] = sum of weights for query q + sum_w = torch.zeros(Q, device=device, dtype=dtype) + sum_w.scatter_add_(0, q_idx_per_nbr, w) + mask = sum_w > 0 + out[mask] = out[mask] / sum_w[mask].unsqueeze(-1) + + return out, input_pos_hash_table + + +class LocalFeatureExtractor(torch.nn.Module): + + def __init__( + self, + # ============= Continuous Convolution Parameters ============= + kernel_size=[4, 4, 4], + cconv_embedding_dim=32, + data=None, + use_topology_conv=False, + topology_search_radius=None, + radius_scale=1.5, + coordinate_mapping="ball_to_cube_volume_preserving", + interpolation="linear", + use_window=True, + particle_radius=0.025, + other_feats_channels=0, + obstacle_feats_channels=3, + topology_blend_alpha=None, + ): + super().__init__() + + # ============= Continuous Convolution Configuration ============= + self.data = data + self.kernel_size = kernel_size + self.cconv_embedding_dim = cconv_embedding_dim + self.radius_scale = radius_scale + self.coordinate_mapping = coordinate_mapping + self.interpolation = interpolation + self.use_window = use_window + self.particle_radius = particle_radius + self.other_feats_channels = other_feats_channels + self.obstacle_feats_channels = obstacle_feats_channels + self.is_local_fluid = self.data == "local_fluid" + self.topology_blend_alpha = topology_blend_alpha + self.filter_extent = np.float32(self.radius_scale * 6 * self.particle_radius) + self.use_topology_conv = bool(use_topology_conv) + self.topology_search_radius = ( + None if topology_search_radius is None else float(topology_search_radius) + ) + self.topology_alpha_param = None + self.topology_alpha_fixed = None + self.topology_fusion_mode = "disabled" + if self.use_topology_conv: + if topology_blend_alpha is None: + self.topology_fusion_mode = "concat" + else: + alpha_fixed = float(topology_blend_alpha) + if alpha_fixed < 0.0 or alpha_fixed > 1.0: + raise ValueError( + f"topology_blend_alpha must be in [0, 1], got {alpha_fixed}" + ) + if alpha_fixed not in (0.0, 1.0): + raise ValueError( + "topology_blend_alpha now only supports null/0/1: " + "null=concat(auto), 0=topology-only, 1=fluid-only." + ) + self.topology_alpha_fixed = alpha_fixed + if alpha_fixed == 0.0: + self.topology_fusion_mode = "topology_only" + else: + self.topology_fusion_mode = "fluid_only" + if self.topology_search_radius is None: + self.topology_filter_extent = self.filter_extent + else: + self.topology_filter_extent = np.float32( + self.radius_scale * 6 * self.topology_search_radius + ) + + # ============= Initialize Neighbor Search ============= + # Define window function for continuous convolution + def window_poly6(r_sqr): + return torch.clamp((1 - r_sqr) ** 3, 0, 1) + + window_fn = window_poly6 if self.use_window else None + + # Initialize particle radius search for neighbor finding + self.particle_radius_search = ParticleRadiusResearch( + radius_search_metric="L2", + radius_search_ignore_query_points=True, + window_function=window_fn, + ) + + # ============= Initialize Continuous Convolution Layers ============= + # Helper function to create continuous convolution layers + def Conv(activation=None, **kwargs): + conv = ContinuousConv( + kernel_size=self.kernel_size, + activation=activation, + align_corners=True, + interpolation=self.interpolation, + coordinate_mapping=self.coordinate_mapping, + normalize=False, + **kwargs, + ) + return conv + + # Fluid-fluid convolution: particle-particle interactions + fluid_conv_out_channels = cconv_embedding_dim + topology_conv_out_channels = cconv_embedding_dim + if self.topology_fusion_mode == "concat": + if cconv_embedding_dim % 2 != 0: + raise ValueError( + f"cconv_embedding_dim must be even for topology concat mode, got {cconv_embedding_dim}" + ) + fluid_conv_out_channels = cconv_embedding_dim // 2 + topology_conv_out_channels = cconv_embedding_dim // 2 + + self.conv0_fluid = Conv( + in_channels=4 + other_feats_channels, + filters=fluid_conv_out_channels, + activation=None, + ) + if self.use_topology_conv: + self.conv0_topology = Conv( + in_channels=4 + other_feats_channels, + filters=topology_conv_out_channels, + activation=None, + ) + + # Fluid-obstacle convolution: particle-boundary interactions + if uses_obstacle_features(self.data): + self.conv0_obstacle = Conv( + in_channels=obstacle_feats_channels, + filters=cconv_embedding_dim, + activation=None, + ) + + # Dense layer for direct feature transformation + dense0_in_features = ( + 7 + other_feats_channels + if self.is_local_fluid + else 4 + other_feats_channels + ) + dense0_out_features = ( + 3 * cconv_embedding_dim if self.is_local_fluid else cconv_embedding_dim + ) + self.dense0_fluid = nn.Linear( + in_features=dense0_in_features, out_features=dense0_out_features + ) + nn.init.xavier_uniform_(self.dense0_fluid.weight) + nn.init.zeros_(self.dense0_fluid.bias) + + self.register_buffer( + "filter_extent_tensor", + torch.tensor(self.filter_extent, dtype=torch.float32), + ) + self.register_buffer( + "topology_filter_extent_tensor", + torch.tensor(self.topology_filter_extent, dtype=torch.float32), + ) + + def forward( + self, + inp_pos, + query_pos, + inp_vel, + other_feats, + box=None, + box_feats=None, + query_vel=None, + fixed_radius_search_hash_table=None, + batch_info=None, + topology_neighbors_index=None, + topology_neighbors_row_splits=None, + ): + filter_extent = self.filter_extent_tensor + topology_filter_extent = self.topology_filter_extent_tensor + + if self.is_local_fluid: + fluid_feats = [torch.ones_like(inp_pos[:, 0:1]), inp_vel, inp_pos] + if other_feats is not None: + if other_feats.dim() == 1: + other_feats = other_feats.unsqueeze(-1) + fluid_feats.append(other_feats) + fluid_feats = torch.cat(fluid_feats, axis=-1) + dense0_fluid_out = self.dense0_fluid(fluid_feats) + particle_features = dense0_fluid_out + return particle_features, None, None + + # ============= Step 1: Neighbor Search ============= + # Find fluid-fluid neighbors (particle-particle interactions) + if batch_info is not None: + batch_size = batch_info["batch_size"] + batch_splits = ( + batch_info["batch_splits"] + .to(device=query_pos.device, dtype=torch.int64) + .contiguous() + ) + if batch_splits.numel() != batch_size + 1: + raise ValueError( + f"batch_splits must contain {batch_size + 1} entries, " + f"got {batch_splits.numel()}" + ) + + fluid_fluid_nns = self.particle_radius_search( + inp_positions=inp_pos, + out_positions=query_pos, + extents=filter_extent, + fixed_radius_search_hash_table=fixed_radius_search_hash_table, + inp_row_splits=batch_splits, + out_row_splits=batch_splits, + ) + + if uses_obstacle_features(self.data): + box_splits = ( + batch_info["box_splits"] + .to(device=box.device, dtype=torch.int64) + .contiguous() + ) + if box_splits.numel() != batch_size + 1: + raise ValueError( + f"box_splits must contain {batch_size + 1} entries, " + f"got {box_splits.numel()}" + ) + + fluid_obstacle_nns = self.particle_radius_search( + inp_positions=box, + out_positions=query_pos, + extents=filter_extent, + fixed_radius_search_hash_table=fixed_radius_search_hash_table, + inp_row_splits=box_splits, + out_row_splits=batch_splits, + ) + + else: + fluid_fluid_nns = self.particle_radius_search( + inp_positions=inp_pos, + out_positions=query_pos, + extents=filter_extent, + fixed_radius_search_hash_table=fixed_radius_search_hash_table, + ) + if uses_obstacle_features(self.data): + # Find fluid-obstacle neighbors (particle-boundary interactions) + fluid_obstacle_nns = self.particle_radius_search( + inp_positions=box, + out_positions=query_pos, + extents=filter_extent, + fixed_radius_search_hash_table=fixed_radius_search_hash_table, + ) + + # ============= Step 2: Feature Extraction via Continuous Convolutions ============= + # Prepare fluid particle features: [1, velocities, other_features] + fluid_feats = [ + torch.ones_like(inp_pos[:, 0:1]), + inp_vel, + ] # Add constant feature for bias + if other_feats is not None: + if other_feats.dim() == 1: + other_feats = other_feats.unsqueeze(-1) + fluid_feats.append(other_feats) + fluid_feats = torch.cat(fluid_feats, axis=-1) + + # Extract fluid branch first when required by fusion mode. + conv0_fluid_out = None + if self.topology_fusion_mode in {"disabled", "concat", "fluid_only"}: + conv0_fluid_out = self.conv0_fluid( + fluid_feats, + inp_pos, + query_pos, + filter_extent, + inp_importance=fluid_fluid_nns["inp_importance"], + neighbors_index=fluid_fluid_nns["neighbors_index"], + neighbors_row_splits=fluid_fluid_nns["neighbors_row_splits"], + neighbors_importance=fluid_fluid_nns["neighbors_importance"], + ) + + topology_neighbors_available = ( + self.use_topology_conv + and topology_neighbors_index is not None + and topology_neighbors_row_splits is not None + ) + + conv0_topology_out = None + # Topology neighborhood branch (enabled when topology neighbors are provided). + if topology_neighbors_available and self.topology_fusion_mode in { + "concat", + "topology_only", + }: + topology_neighbors_index = topology_neighbors_index.to( + device=inp_pos.device, dtype=torch.int32 + ).contiguous() + topology_neighbors_row_splits = topology_neighbors_row_splits.to( + device=inp_pos.device, dtype=torch.int64 + ).contiguous() + expected_num_row_splits = query_pos.shape[0] + 1 + if topology_neighbors_row_splits.numel() != expected_num_row_splits: + raise ValueError( + f"topology_neighbors_row_splits size mismatch: got " + f"{topology_neighbors_row_splits.numel()}, expected {expected_num_row_splits}" + ) + + topology_nns = self.particle_radius_search( + inp_positions=inp_pos, + out_positions=query_pos, + extents=topology_filter_extent, + user_neighbors_index=topology_neighbors_index, + user_neighbors_row_splits=topology_neighbors_row_splits, + user_neighbors_importance=torch.ones_like( + topology_neighbors_index, dtype=torch.float32, device=inp_pos.device + ), + ) + conv0_topology_out = self.conv0_topology( + fluid_feats, + inp_pos, + query_pos, + topology_filter_extent, + inp_importance=topology_nns["inp_importance"], + neighbors_index=topology_nns["neighbors_index"], + neighbors_row_splits=topology_nns["neighbors_row_splits"], + neighbors_importance=topology_nns["neighbors_importance"], + ) + + # Topology + fluid fusion policy: + # - disabled: fluid-only + # - concat: [fluid_half, topology_half] -> cconv_embedding_dim + # - topology_only: use topology branch only (alpha=0) + # - fluid_only: use fluid branch only (alpha=1) + if self.topology_fusion_mode == "concat": + if conv0_fluid_out is None: + raise ValueError("concat mode expects fluid branch output, got None.") + if conv0_topology_out is None: + conv0_topology_out = torch.zeros_like(conv0_fluid_out) + conv0_fluid_out = torch.cat([conv0_fluid_out, conv0_topology_out], dim=-1) + elif self.topology_fusion_mode == "topology_only": + if conv0_topology_out is None: + raise ValueError( + "topology_blend_alpha=0 requires topology neighbors, " + "but topology neighbors are missing in the current batch." + ) + conv0_fluid_out = conv0_topology_out + elif self.topology_fusion_mode == "fluid_only": + if conv0_fluid_out is None: + raise ValueError( + "topology_blend_alpha=1 expects fluid branch output, got None." + ) + elif self.topology_fusion_mode == "disabled": + if conv0_fluid_out is None: + raise ValueError( + "topology disabled mode expects fluid branch output, got None." + ) + else: + raise ValueError( + f"Unsupported topology_fusion_mode={self.topology_fusion_mode}" + ) + + # Dense layer: Direct transformation of particle features + if query_vel is not None: + query_fluid_feats = [torch.ones_like(query_pos[:, 0:1]), query_vel] + dense0_fluid_out = self.dense0_fluid(torch.cat(query_fluid_feats, axis=-1)) + else: + dense0_fluid_out = self.dense0_fluid(fluid_feats) + + if uses_obstacle_features(self.data): + if box_feats is None: + raise ValueError("box_feats is required for obstacle-feature data.") + if box_feats.dim() == 1: + box_feats = box_feats.unsqueeze(-1) + if box_feats.shape[-1] != self.obstacle_feats_channels: + raise ValueError( + f"box_feats channel mismatch: got {box_feats.shape[-1]}, " + f"expected {self.obstacle_feats_channels}. " + "Set `obstacle_feats_channels` to match your data." + ) + # Fluid-obstacle convolution: Extract features from boundary interactions + conv0_obstacle_out = self.conv0_obstacle( + box_feats, + box, + query_pos, + filter_extent, + inp_importance=fluid_obstacle_nns["inp_importance"], + neighbors_index=fluid_obstacle_nns["neighbors_index"], + neighbors_row_splits=fluid_obstacle_nns["neighbors_row_splits"], + neighbors_importance=fluid_obstacle_nns["neighbors_importance"], + ) + + # Combine all extracted features + particle_features = torch.cat( + [ + conv0_obstacle_out, # Boundary interaction features + conv0_fluid_out, # Particle-particle features + dense0_fluid_out, # Direct particle features + ], + axis=-1, + ) # Shape: [N, combined_feat_dim] + + else: + particle_features = torch.cat( + [ + conv0_fluid_out, # Particle-particle features + dense0_fluid_out, # Direct particle features + ], + axis=-1, + ) + + num_fluid_neighbors = reduce_subarrays_sum( + torch.ones_like(fluid_fluid_nns["neighbors_index"], dtype=torch.float32), + fluid_fluid_nns["neighbors_row_splits"], + ) + return particle_features, num_fluid_neighbors, fluid_fluid_nns + + +class SuperParticleLocalFeatureExtractor(torch.nn.Module): + + def __init__( + self, + super_particle_compression_ratio=10.0, + # ============= Continuous Convolution Parameters ============= + kernel_size=[4, 4, 4], + cconv_embedding_dim=32, + radius_scale=1.5, + coordinate_mapping="ball_to_cube_volume_preserving", + interpolation="linear", + use_window=True, + particle_radius=0.025, + other_feats_channels=0, + ): + super().__init__() + + # ============= Continuous Convolution Configuration ============= + self.kernel_size = kernel_size + self.cconv_embedding_dim = cconv_embedding_dim + self.radius_scale = radius_scale + self.coordinate_mapping = coordinate_mapping + self.interpolation = interpolation + self.use_window = use_window + self.particle_radius = particle_radius + self.other_feats_channels = other_feats_channels + self.filter_extent = np.float32( + self.radius_scale + * 6 + * self.particle_radius + * math.pow(super_particle_compression_ratio, 1.0 / 3.0) + ) + + # ============= Initialize Neighbor Search ============= + # Define window function for continuous convolution + def window_poly6(r_sqr): + return torch.clamp((1 - r_sqr) ** 3, 0, 1) + + window_fn = window_poly6 if self.use_window else None + + # Initialize particle radius search for neighbor finding + self.particle_radius_search = ParticleRadiusResearch( + radius_search_metric="L2", + radius_search_ignore_query_points=True, + window_function=window_fn, + ) + + # ============= Initialize Continuous Convolution Layers ============= + # Helper function to create continuous convolution layers + def Conv(activation=None, **kwargs): + conv = ContinuousConv( + kernel_size=self.kernel_size, + activation=activation, + align_corners=True, + interpolation=self.interpolation, + coordinate_mapping=self.coordinate_mapping, + normalize=False, + **kwargs, + ) + return conv + + # Fluid-fluid convolution: particle-particle interactions + self.conv0_fluid = Conv( + in_channels=4 + other_feats_channels, + filters=cconv_embedding_dim, + activation=None, + ) + + # Fluid-obstacle convolution: particle-boundary interactions + self.conv0_obstacle = Conv( + in_channels=3, filters=cconv_embedding_dim, activation=None + ) + + # Dense layer for direct feature transformation + self.dense0_fluid = nn.Linear( + in_features=4 + other_feats_channels, out_features=cconv_embedding_dim + ) + nn.init.xavier_uniform_(self.dense0_fluid.weight) + nn.init.zeros_(self.dense0_fluid.bias) + + def forward( + self, + inp_pos, + query_pos, + inp_vel, + other_feats, + box, + box_feats, + query_vel=None, + inp_pos_hash_table=None, + box_hash_table=None, + ): + filter_extent = torch.tensor( + self.filter_extent, device=inp_pos.device, dtype=inp_pos.dtype + ) + # ============= Step 1: Neighbor Search ============= + # Find fluid-fluid neighbors (particle-particle interactions) + if inp_pos_hash_table is not None: + fluid_fluid_nns = self.particle_radius_search( + inp_positions=inp_pos, + out_positions=query_pos, + extents=filter_extent, + fixed_radius_search_hash_table=inp_pos_hash_table, + ) + else: + fluid_fluid_nns, inp_pos_hash_table = self.particle_radius_search( + inp_positions=inp_pos, + out_positions=query_pos, + extents=filter_extent, + return_hash_table=True, + ) + + # Find fluid-obstacle neighbors (particle-boundary interactions) + if box_hash_table is not None: + fluid_obstacle_nns = self.particle_radius_search( + inp_positions=box, + out_positions=query_pos, + extents=filter_extent, + fixed_radius_search_hash_table=box_hash_table, + ) + else: + fluid_obstacle_nns, box_hash_table = self.particle_radius_search( + inp_positions=box, + out_positions=query_pos, + extents=filter_extent, + return_hash_table=True, + ) + + # ============= Step 2: Feature Extraction via Continuous Convolutions ============= + # Prepare fluid particle features: [1, velocities, other_features] + fluid_feats = [ + torch.ones_like(inp_pos[:, 0:1]), + inp_vel, + ] # Add constant feature for bias + if other_feats is not None: + fluid_feats.append(other_feats) + fluid_feats = torch.cat(fluid_feats, axis=-1) + + # Extract features using continuous convolutions + # Fluid-fluid convolution: Extract features from neighboring particles + conv0_fluid_out = self.conv0_fluid( + fluid_feats, + inp_pos, + query_pos, + filter_extent, + inp_importance=fluid_fluid_nns["inp_importance"], + neighbors_index=fluid_fluid_nns["neighbors_index"], + neighbors_row_splits=fluid_fluid_nns["neighbors_row_splits"], + neighbors_importance=fluid_fluid_nns["neighbors_importance"], + ) + + # Dense layer: Direct transformation of particle features + if query_vel is not None: + query_fluid_feats = [torch.ones_like(query_pos[:, 0:1]), query_vel] + dense0_fluid_out = self.dense0_fluid(torch.cat(query_fluid_feats, axis=-1)) + else: + dense0_fluid_out = self.dense0_fluid(fluid_feats) + + # Fluid-obstacle convolution: Extract features from boundary interactions + conv0_obstacle_out = self.conv0_obstacle( + box_feats, + box, + query_pos, + filter_extent, + inp_importance=fluid_obstacle_nns["inp_importance"], + neighbors_index=fluid_obstacle_nns["neighbors_index"], + neighbors_row_splits=fluid_obstacle_nns["neighbors_row_splits"], + neighbors_importance=fluid_obstacle_nns["neighbors_importance"], + ) + + # Combine all extracted features + particle_features = torch.cat( + [ + conv0_obstacle_out, # Boundary interaction features + conv0_fluid_out, # Particle-particle features + dense0_fluid_out, # Direct particle features + ], + axis=-1, + ) # Shape: [N, combined_feat_dim] + + return particle_features, inp_pos_hash_table, box_hash_table + + +class SuperParticleDecoderLayer(nn.Module): + """ + A single decoder layer that performs: + 1. Cross-attention from query particles to super-particles + 2. Self-attention among query particles + 3. Feed-forward network + + This follows a standard transformer decoder architecture with pre-normalization. + """ + + def __init__( + self, + query_particle_feat_dim: int, + super_particle_feat_dim: int, + num_heads: int = 8, + ffn_hidden_dim: int = 256, + dropout: float = 0.1, + bias: bool = True, + activation: str = "swiglu", + norm_type: str = "rms_norm", + qk_norm: bool = True, + enable_self_attn: bool = True, + enable_ffn: bool = True, + ): + """ + Args: + query_particle_feat_dim: Feature dimension of query particles + super_particle_feat_dim: Feature dimension of super-particles (keys/values) + num_heads: Number of attention heads + ffn_hidden_dim: Hidden dimension for feed-forward network + dropout: Dropout rate + bias: Whether to use bias in linear layers + activation: FFN activation function ('swiglu' or 'gelu') + norm_type: Normalization type ('layer_norm' or 'rms_norm') + qk_norm: Whether to apply QK normalization in attention + enable_self_attn: Whether to keep self-attention in this layer + enable_ffn: Whether to keep FFN in this layer + """ + super().__init__() + + self.query_particle_feat_dim = query_particle_feat_dim + self.super_particle_feat_dim = super_particle_feat_dim + self.num_heads = num_heads + self.head_dim = query_particle_feat_dim // num_heads + self.enable_self_attn = bool(enable_self_attn) + self.enable_ffn = bool(enable_ffn) + + assert ( + query_particle_feat_dim % num_heads == 0 + ), f"query_particle_feat_dim {query_particle_feat_dim} must be divisible by num_heads {num_heads}" + + # Choose normalization module + if norm_type == "layer_norm": + norm_module = nn.LayerNorm + elif norm_type == "rms_norm": + norm_module = RMSNorm + else: + raise ValueError(f"Unsupported norm_type: {norm_type}") + + # Dropout for residual connections + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + + # ============= Cross-Attention Components ============= + # Pre-norm for cross-attention + self.query_norm = norm_module(query_particle_feat_dim, eps=1e-6) + self.cross_kv_norm = norm_module(super_particle_feat_dim, eps=1e-6) + + # Multi-head cross-attention: query particles attend to super-particles + self.multihead_cross_attn = MultiHeadAttention( + query_dim=query_particle_feat_dim, + num_heads=num_heads, + kv_dim=super_particle_feat_dim, + bias=bias, + qk_norm=qk_norm, + norm_type=norm_type, + ) + + # ============= Self-Attention Components ============= + if self.enable_self_attn: + # Pre-norm for self-attention + self.self_attn_norm = norm_module(query_particle_feat_dim, eps=1e-6) + + # Multi-head self-attention among query particles + self.self_attn = MultiHeadAttention( + query_dim=query_particle_feat_dim, + num_heads=num_heads, + kv_dim=None, # Self-attention + bias=bias, + qk_norm=qk_norm, + norm_type=norm_type, + ) + else: + self.self_attn_norm = None + self.self_attn = None + + # ============= Feed-Forward Network ============= + if self.enable_ffn: + # Pre-norm for FFN + self.ffn_norm = norm_module(query_particle_feat_dim, eps=1e-6) + + # Choose FFN activation type + if activation == "swiglu": + self.ffn = FeedForwardSwiGLU( + dim=query_particle_feat_dim, + hidden_dim=ffn_hidden_dim, + dropout=dropout, + bias=bias, + ) + elif activation == "gelu": + self.ffn = FeedForwardGeLU( + dim=query_particle_feat_dim, + hidden_dim=ffn_hidden_dim, + dropout=dropout, + bias=bias, + ) + else: + raise ValueError(f"Unsupported activation: {activation}") + else: + self.ffn_norm = None + self.ffn = None + + def forward( + self, + query: torch.Tensor, + kv: torch.Tensor, + rope_cos: torch.Tensor = None, + rope_sin: torch.Tensor = None, + rope_ctx_cos: torch.Tensor = None, + rope_ctx_sin: torch.Tensor = None, + ): + """ + Forward pass of the decoder layer. + + Args: + query: Query particle features [B, N_query, query_dim] + kv: Super-particle features (keys/values) [B, N_super, super_dim] + rope_cos: Cosine RoPE embeddings for queries + rope_sin: Sine RoPE embeddings for queries + rope_ctx_cos: Cosine RoPE embeddings for keys/values + rope_ctx_sin: Sine RoPE embeddings for keys/values + + Returns: + Updated query particle features [B, N_query, query_dim] + """ + # ============= Cross-Attention Block ============= + # Query particles attend to super-particles + residual = query + q = self.query_norm(query) + kv_normed = self.cross_kv_norm(kv) + + cross_attn_output = self.multihead_cross_attn( + q=q, + k=kv_normed, + v=kv_normed, + src_key_padding_mask=None, + rope_cos=rope_cos, + rope_sin=rope_sin, + rope_ctx_cos=rope_ctx_cos, + rope_ctx_sin=rope_ctx_sin, + ) + + query = residual + self.dropout(cross_attn_output) + + # ============= Self-Attention Block ============= + # Query particles attend to each other + if self.enable_self_attn: + residual = query + q = self.self_attn_norm(query) + + self_attn_output = self.self_attn( + q=q, + k=q, + v=q, + src_key_padding_mask=None, + rope_cos=rope_cos, + rope_sin=rope_sin, + rope_ctx_cos=None, + rope_ctx_sin=None, + ) + query = residual + self.dropout(self_attn_output) + + # ============= Feed-Forward Block ============= + if self.enable_ffn: + residual = query + ffn_output = self.ffn(self.ffn_norm(query)) + query = residual + self.dropout(ffn_output) + + return query + + +class SuperParticleDecoder(nn.Module): + """ + Super-Particle Decoder that stacks self-attention, cross-attention, and FFN layers. + + This decoder processes super-particles (from slot attention) by: + 1. Self-attention among super-particles (with optional RoPE) + 2. Cross-attention from super-particles to all particles (with RoPE) + 3. Feed-forward network + + All with residual connections and layer normalization. + + Architecture inspired by standard transformer decoder blocks. + + """ + + def __init__( + self, + num_layers: int, + query_particle_feat_dim: int, + super_particle_feat_dim: int, + num_heads: int = 8, + ffn_hidden_dim: int = 256, + dropout: float = 0.1, + bias: bool = True, + activation: str = "swiglu", + norm_type: str = "rms_norm", + qk_norm: bool = True, + enable_self_attn: bool = True, + enable_ffn: bool = True, + ): + """ + Args: + dim: Model dimension (must match input feature dim) + num_heads: Number of attention heads + ffn_hidden_dim: Hidden dimension for FFN + dropout: Dropout rate + bias: Whether to use bias in linear layers + activation: FFN activation ('swiglu' or 'gelu') + norm_type: Normalization type ('layer_norm' or 'rms_norm') + qk_norm: Whether to apply QK normalization in attention + enable_self_attn: Whether to keep self-attention in each decoder layer + enable_ffn: Whether to keep FFN in each decoder layer + use_rope_self_attn: Whether to use RoPE in self-attention + use_rope_cross_attn: Whether to use RoPE in cross-attention + rope_dim: Dimension for RoPE encoding (auto-computed if None) + """ + super().__init__() + + self.query_particle_feat_dim = query_particle_feat_dim + self.super_particle_feat_dim = super_particle_feat_dim + self.num_heads = num_heads + self.head_dim = query_particle_feat_dim // num_heads + assert ( + query_particle_feat_dim % num_heads == 0 + ), f"query_particle_feat_dim {query_particle_feat_dim} must be divisible by num_heads {num_heads}" + + # layers + self.layers = nn.ModuleList( + [ + SuperParticleDecoderLayer( + query_particle_feat_dim=query_particle_feat_dim, + super_particle_feat_dim=super_particle_feat_dim, + num_heads=num_heads, + ffn_hidden_dim=ffn_hidden_dim, + dropout=dropout, + bias=bias, + activation=activation, + norm_type=norm_type, + qk_norm=qk_norm, + enable_self_attn=enable_self_attn, + enable_ffn=enable_ffn, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + super_particles_feat: torch.Tensor, + particles_feat: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + rope_ctx_cos: torch.Tensor, + rope_ctx_sin: torch.Tensor, + ): + for idx, layer in enumerate(self.layers): + particles_feat = layer( + query=particles_feat, + kv=super_particles_feat, + rope_cos=rope_cos, + rope_sin=rope_sin, + rope_ctx_cos=rope_ctx_cos, + rope_ctx_sin=rope_ctx_sin, + ) + return particles_feat diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_runtime.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_runtime.py new file mode 100644 index 000000000..3ed5f3d0f --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/_runtime.py @@ -0,0 +1,249 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Checkpoint-compatible WorldParticle inference, independent of live state.""" + +from __future__ import annotations +from typing import Any +from pathlib import Path +import math +import warnings + +import numpy as np + +SUPPORTED_TASKS = ("fancy_fluid", "cloth_on_ball", "duck_drop", "fluid") + + +def load_config(path: str | Path | None) -> dict[str, Any]: + """Read one of the four supported official model configurations.""" + import yaml + + if path is None or not Path(path).is_file(): + raise ValueError("WorldParticle config must name an existing local file") + with Path(path).open() as stream: + cfg = yaml.safe_load(stream) + if not isinstance(cfg, dict) or cfg.get("data") not in SUPPORTED_TASKS: + raise ValueError(f"WorldParticle supported data tasks are {SUPPORTED_TASKS}") + dt = float(cfg.get("timestep", 0.0)) + if not math.isfinite(dt) or dt <= 0: + raise ValueError("Checkpoint timestep must be positive and finite") + if cfg.get("verbose", False): + raise ValueError("WorldParticle verbose/debug mode must be disabled") + gravity = np.asarray(cfg.get("gravity"), dtype=np.float32) + if gravity.shape != (3,) or not np.isfinite(gravity).all(): + raise ValueError("Checkpoint gravity must be a finite 3-vector") + return dict(cfg, timestep=dt) + + +def normalize_features(task: str, raw: np.ndarray) -> np.ndarray: + """Preserve each official dataset's feature transform and input precision.""" + channels = {"fancy_fluid": 3, "cloth_on_ball": 7, "duck_drop": 1, "fluid": 2} + if raw.ndim != 2 or raw.shape[1] != channels[task] or not len(raw): + raise ValueError(f"{task} raw feats must have shape (N, {channels[task]})") + if not np.issubdtype(raw.dtype, np.number) or not np.isfinite(raw).all(): + raise ValueError("WorldParticle features must be finite numeric values") + if task == "fancy_fluid": + feature = raw[:, 2:3] + feature = (feature - feature.mean(axis=0)) / (feature.std(axis=0) + 1e-8) + elif task == "cloth_on_ball": + mass, young, density = raw[:, 0], raw[:, -2], raw[:, -1] + if np.any(young <= 0): + raise ValueError("Cloth Young modulus feature must be positive") + young = (np.log10(young) - np.log10(5e3)) / (np.log10(3e5) - np.log10(5e3)) + density = (density - 50.0) / (500.0 - 50.0) + mass = (mass - mass.mean(axis=0)) / (mass.std(axis=0) + 1e-8) + feature = np.concatenate( + [mass[:, None], young[:, None], density[:, None]], axis=1 + ) + else: + feature = raw + feature = np.asarray(feature, dtype=np.float32) + if not np.isfinite(feature).all(): + raise ValueError("Normalized WorldParticle features must be finite") + return feature + + +def load_assets(options: dict[str, Any]) -> tuple: + """Load static NPZ features/boundaries/topology; never trajectory state. + + Config task selects raw feature semantics. Topology-enabled checkpoints + require the original CSR neighbors, not a newly inferred physical graph. + """ + from ._network._data import ( + uses_topology_neighbors, + get_default_topology_search_radius, + ) + + unknown = set(options) - {"checkpoint", "config", "inputs", "model_to_world"} + if unknown: + raise ValueError(f"Unknown WorldParticle options: {sorted(unknown)}") + for key in ("checkpoint", "config", "inputs"): + if key not in options or not Path(options[key]).is_file(): + raise ValueError(f"WorldParticle {key} must name an existing local file") + cfg = load_config(options["config"]) + task = cfg["data"] + use_topology = cfg.get("use_topology_conv") + if use_topology is None: + use_topology = uses_topology_neighbors(task) + topology_radius = cfg.get("topology_search_radius") + if topology_radius is None and use_topology: + topology_radius = get_default_topology_search_radius(task) + rotation = np.asarray(options.get("model_to_world", np.eye(3)), dtype=np.float32) + if ( + rotation.shape != (3, 3) + or not np.isfinite(rotation).all() + or not np.allclose(rotation @ rotation.T, np.eye(3), atol=1e-6) + or not np.isclose(np.linalg.det(rotation), 1.0, atol=1e-6) + ): + raise ValueError("model_to_world must be a proper orthonormal 3x3 rotation") + with np.load(options["inputs"], allow_pickle=False) as data: + raw = data["feats"] + feature = normalize_features(task, raw) + box = np.asarray(data["box"], dtype=np.float32) + box_feats = np.asarray(data["box_feats"], dtype=np.float32) + topology = () + if use_topology: + if not {"neighbors_index", "neighbors_row_splits"}.issubset(data.files): + raise ValueError( + "Topology-enabled checkpoint requires CSR neighbors_index/row_splits" + ) + index = data["neighbors_index"] + rows = data["neighbors_row_splits"] + if ( + index.ndim != 1 + or rows.ndim != 1 + or len(rows) != len(feature) + 1 + or not np.issubdtype(index.dtype, np.integer) + or not np.issubdtype(rows.dtype, np.integer) + or rows[0] != 0 + or rows[-1] != len(index) + or np.any(np.diff(rows) < 0) + or np.any(index < 0) + or np.any(index >= len(feature)) + ): + raise ValueError("Invalid particle topology CSR layout or index range") + topology = (index.astype(np.int64), rows.astype(np.int64)) + if ( + box.ndim != 2 + or box.shape[1] != 3 + or not len(box) + or box_feats.shape != box.shape + ): + raise ValueError("box and box_feats must have nonempty matching shapes (M, 3)") + if not all(np.isfinite(x).all() for x in (box, box_feats)): + raise ValueError("WorldParticle static inputs must be finite") + cfg.update( + resolved_other_feats_channels=feature.shape[1], + resolved_obstacle_feats_channels=3, + resolved_use_topology_conv=bool(use_topology), + resolved_topology_search_radius=topology_radius, + ) + return cfg, (feature, box, box_feats) + topology, rotation + + +def load_checkpoint_state( + network: Any, state: dict[str, Any], *, legacy_fluid: bool = False +) -> None: + """Strictly load weights, migrating only one known unused legacy buffer.""" + state = dict(state) + missing = set(network.state_dict()) - set(state) + legacy = "local_feature_extractor.topology_filter_extent_tensor" + if legacy_fluid and missing == {legacy}: + state[legacy] = network.state_dict()[legacy] + warnings.warn( + "Legacy Deep Fluid checkpoint omits the unused topology_filter_extent_tensor buffer; using its configured default. All trainable weights remain strict.", + UserWarning, + stacklevel=2, + ) + network.load_state_dict(state, strict=True) + + +def load_network(cfg: dict[str, Any], checkpoint: str | Path, device: str) -> Any: + """Construct the original architecture and strictly load its model tensors.""" + import sys + import platform + import torch + from importlib.metadata import version + + if ( + sys.platform != "linux" + or platform.machine() != "x86_64" + or sys.version_info[:2] != (3, 12) + ): + raise RuntimeError( + "WorldParticle currently requires Linux x86_64 / Python 3.12" + ) + if ( + torch.__version__ != "2.7.1+cu128" + or not torch._C._GLIBCXX_USE_CXX11_ABI + or torch.cuda.get_device_capability(device) != (12, 0) + ): + raise RuntimeError( + "WorldParticle requires the validated Torch 2.7.1+cu128 / CXX11 ABI / SM120 runtime" + ) + if ( + version("open3d") != "0.19.0+1e7b174" + or version("flash-attn").split("+")[0] != "2.8.3" + ): + raise RuntimeError( + "Install EmbodiChain's validated worldparticle extra; the Open3D/FlashAttention operators do not match" + ) + from ._network.particle_network_cross_attn_feat import ( + ParticleNetworkCrossAttnLocalFeature, + ) + + names = ( + "kernel_size", + "cconv_embedding_dim", + "radius_scale", + "particle_radius", + "slot_compression_ratio", + "slot_attn_iters", + "slot_attn_heads", + "particle_position_rope_dim", + "num_decoder_layers", + "decoder_attn_heads", + "decoder_attn_ffn_hidden_dim", + "decoder_attn_dropout", + "output_hidden_dim", + "output_layers", + "velhead", + "topology_blend_alpha", + ) + kwargs = {key: cfg[key] for key in names if key in cfg} + kwargs.update( + data=cfg["data"], + gravity=cfg["gravity"], + timestep=cfg["timestep"], + other_feats_channels=cfg["resolved_other_feats_channels"], + obstacle_feats_channels=cfg["resolved_obstacle_feats_channels"], + verbose=False, + use_topology_conv=cfg["resolved_use_topology_conv"], + topology_search_radius=cfg["resolved_topology_search_radius"], + ) + network = ParticleNetworkCrossAttnLocalFeature(**kwargs) + checkpoint_data = torch.load(checkpoint, map_location="cpu", weights_only=True) + state = checkpoint_data["state_dict"] + prefix = "model." + if not state or any(not key.startswith(prefix) for key in state): + raise ValueError( + "Expected a WorldParticle Lightning checkpoint with model.* tensors" + ) + load_checkpoint_state( + network, + {key[len(prefix) :]: value for key, value in state.items()}, + legacy_fluid=cfg["data"] == "fluid" and not cfg["resolved_use_topology_conv"], + ) + return network.to(device).eval().requires_grad_(False) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/data.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/data.py new file mode 100644 index 000000000..2a7bad766 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/data.py @@ -0,0 +1,90 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Official NPZ initial-state slicing and authored surface topology.""" + +from __future__ import annotations +from itertools import combinations +from pathlib import Path +from typing import Any +import numpy as np + +__all__ = ["load_sample"] + + +def cloth_faces(points, indices, rows): + """Recover the display triangles from the sample's four-neighbor grid.""" + adjacency = [ + set(map(int, indices[rows[i] : rows[i + 1]])) - {i} for i in range(len(points)) + ] + if max(map(len, adjacency), default=0) > 16: + raise ValueError("Cloth sample is not a triangle/four-neighbor surface graph") + faces = set() + for a, neighbors in enumerate(adjacency): + for b in neighbors: + if b > a: + faces.update((a, b, c) for c in neighbors & adjacency[b] if c > b) + if not faces: + quads = set() + for a, neighbors in enumerate(adjacency): + for b, d in combinations(sorted(neighbors), 2): + for c in (adjacency[b] & adjacency[d]) - {a}: + key = tuple(sorted((a, b, c, d))) + if key not in quads: + quads.add(key) + faces.update(((a, b, c), (a, c, d))) + result = np.asarray(sorted(faces), dtype=np.int32).reshape(-1, 3) + if not len(result): + raise ValueError("No cloth display faces recovered from the sample topology") + normals = np.cross( + points[result[:, 1]] - points[result[:, 0]], + points[result[:, 2]] - points[result[:, 0]], + ) + # Omit coincident-point display faces; preserve the original neural CSR. + valid = np.linalg.norm(normals, axis=1) > 1e-12 + result, normals = result[valid], normals[valid] + flip = normals[:, 1] < 0 + result[flip] = result[flip][:, [0, 2, 1]] + return result + + +def load_sample( + path: str | Path, + profile: dict[str, Any], + *, + frame_start: int = 0, + frame_end: int | None = None, +) -> tuple: + """Separate initial dynamics inputs from optional reference-only metrics.""" + with np.load(path, allow_pickle=False) as data: + # Match official inference: task stride, config slice, then initial state. + positions = data["position"][:: profile["stride"]][frame_start:frame_end] + velocities = data["velocity"][:: profile["stride"]][frame_start:frame_end] + if frame_start < 0 or not len(positions): + raise ValueError("Initial frame selection must be nonnegative and nonempty") + q = np.asarray(positions[0], np.float32) + qd = np.asarray(velocities[0], np.float32) + raw = data["feats"].copy() + box = np.asarray(data["box"], np.float32) + normals = np.asarray(data["box_feats"], np.float32) + reference = positions.copy() if profile["reference"] else None + faces = None + if profile["checkpoint"] == "cloth_on_ball": + faces = cloth_faces( + q, data["neighbors_index"], data["neighbors_row_splits"] + ) + if q.shape != qd.shape or q.ndim != 2 or q.shape[1] != 3: + raise ValueError("Initial position/velocity must have matching shapes (N,3)") + return q, qd, raw, box, normals, faces, reference diff --git a/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/model.py b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/model.py new file mode 100644 index 000000000..64a60fe46 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/worldparticle/model.py @@ -0,0 +1,82 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Checkpoint-compatible stateless particle transitions in world coordinates.""" + +from __future__ import annotations +from collections.abc import Mapping +from typing import Any +import torch +from torch import Tensor +from ._runtime import load_assets, load_network + +__all__ = ["WorldParticleModel", "load_model"] + + +class WorldParticleModel: + """Own frozen weights, authored features/topology and model-frame conversion. + + Inputs and outputs use world-space particle order and SI units. Acceleration + includes gravity; model-specific BF16 arithmetic matches official inference. + Callers own execution streams and must create/use this object on that stream. + """ + + def __init__( + self, + network: torch.nn.Module, + inputs: tuple[Tensor, ...], + rotation: Any, + cfg: Mapping[str, Any], + ) -> None: + self.network = network + self.inputs = inputs + self.rotation = torch.as_tensor(rotation, device=inputs[0].device) + self.dt = float(cfg["timestep"]) + self.k_partite = int(cfg.get("k_partite", -1)) + + def __call__( + self, position: Tensor, velocity: Tensor, acceleration: Tensor, dt: float + ) -> tuple[Tensor, Tensor]: + import math + + if not math.isclose(dt, self.dt, rel_tol=1e-7, abs_tol=1e-10): + raise ValueError("WorldParticle step dt must match the checkpoint timestep") + feature, box, box_feats = self.inputs[:3] + arguments = dict( + k_partite=self.k_partite, external_force=acceleration @ self.rotation + ) + if len(self.inputs) == 5: + arguments.update( + topology_neighbors_index=self.inputs[3], + topology_neighbors_row_splits=self.inputs[4], + ) + model_position = position @ self.rotation + model_velocity = velocity @ self.rotation + with torch.autocast("cuda", dtype=torch.bfloat16): + q, qd = self.network( + (model_position, model_velocity, feature, box, box_feats), **arguments + ) + return q.float() @ self.rotation.T, qd.float() @ self.rotation.T + + +def load_model( + options: Mapping[str, Any], device: str = "cuda:0" +) -> WorldParticleModel: + """Load validated static inputs and strict official checkpoint tensors.""" + cfg, arrays, rotation = load_assets(options) + inputs = tuple(torch.as_tensor(value, device=device) for value in arrays) + return WorldParticleModel( + load_network(cfg, options["checkpoint"], device), inputs, rotation, cfg + ) diff --git a/pyproject.toml b/pyproject.toml index 1bc0109d3..4e673d25d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,15 @@ dependencies = [ ] [project.optional-dependencies] +# Only this validated CUDA operator ABI is supported. +worldparticle = [ + "torchvision==0.22.1+cu128; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.12'", + "numpy==2.2.6; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.12'", + "torch==2.7.1+cu128; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.12'", + "einops==0.8.2; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.12'", + "open3d @ http://192.168.3.43:8080/packages/open3d-0.19.0%2B1e7b174-1torch271cu128sm120-cp312-cp312-manylinux_2_35_x86_64.whl#sha256=209bcee758d25953bcc4cab284ad4fb17f5aace0a3e1e3c296750511bf4257ad ; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.12'", + "flash-attn @ http://192.168.3.43:8080/packages/flash_attn-2.8.3%2Bcu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl#sha256=cafd4f708725e1be95495a86cbf418cfada08e707c8a797def61488a2a3c4f72 ; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.12'", +] rigidformer = [ "einops>=0.8.1", "roma==1.5.4", @@ -86,6 +95,9 @@ gensim = [ # decomposition, benchmarks, and workspace cache management. embodichain = "embodichain.__main__:main" +[project.entry-points."dexsim.newton_solvers"] +worldparticle = "embodichain.lab.sim.physics.worldparticle:WorldParticleProvider" + [project.entry-points."embodichain.tasks"] embodichain_tasks = "embodichain_tasks" diff --git a/setup.py b/setup.py index 5e351d06c..897eac207 100644 --- a/setup.py +++ b/setup.py @@ -189,6 +189,10 @@ def main(): package_data={ "embodichain": ["VERSION"], "embodichain_models.dynamics.nerd": ["THIRD_PARTY_LICENSES"], + "embodichain_models.dynamics.worldparticle": [ + "LICENSE.*", + "PROVENANCE.json", + ], "embodichain_models.dynamics.rigidformer": [ "THIRD_PARTY_LICENSES", "REFERENCE_SOURCES.json", diff --git a/tests/models/dynamics/worldparticle/test_data.py b/tests/models/dynamics/worldparticle/test_data.py new file mode 100644 index 000000000..f891c0a9e --- /dev/null +++ b/tests/models/dynamics/worldparticle/test_data.py @@ -0,0 +1,45 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +from __future__ import annotations +import numpy as np +import pytest +from embodichain_models.dynamics.worldparticle.data import load_sample + + +def test_initial_state_and_reference_follow_config_slice(tmp_path): + positions = np.arange(24 * 3 * 3, dtype=np.float32).reshape(24, 3, 3) + velocities = -positions - 1 + path = tmp_path / "sample.npz" + np.savez( + path, + position=positions, + velocity=velocities, + feats=np.ones((3, 3)), + box=np.zeros((1, 3)), + box_feats=np.ones((1, 3)), + ) + profile = dict(checkpoint="fancy_fluid", stride=1, reference=True) + q, qd, *_, reference = load_sample(path, profile, frame_start=10, frame_end=20) + np.testing.assert_array_equal(q, positions[10]) + np.testing.assert_array_equal(qd, velocities[10]) + np.testing.assert_array_equal(reference, positions[10:20]) + profile["stride"] = 5 + q, qd, *_, reference = load_sample(path, profile, frame_start=1, frame_end=4) + np.testing.assert_array_equal(q, positions[5]) + np.testing.assert_array_equal(qd, velocities[5]) + np.testing.assert_array_equal(reference, positions[::5][1:4]) + with pytest.raises(ValueError, match="nonempty"): + load_sample(path, profile, frame_start=10) diff --git a/tests/models/dynamics/worldparticle/test_worldparticle.py b/tests/models/dynamics/worldparticle/test_worldparticle.py new file mode 100644 index 000000000..f91d0d1f7 --- /dev/null +++ b/tests/models/dynamics/worldparticle/test_worldparticle.py @@ -0,0 +1,391 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Optional WorldParticle adapter contracts, without downloading weights.""" + +from __future__ import annotations +from pathlib import Path + +import numpy as np +import pytest +import torch +import warp as wp + +from embodichain_models.dynamics.worldparticle._runtime import load_assets + + +@pytest.fixture +def assets(tmp_path): + cfg = tmp_path / "config.yaml" + cfg.write_text( + "data: fancy_fluid\ntimestep: 0.01\ngravity: [0., 0., 0.]\nk_partite: 2\n" + ) + ckpt = tmp_path / "model.ckpt" + ckpt.write_bytes(b"network loader is replaced only in adapter tests") + data = tmp_path / "inputs.npz" + np.savez( + data, + feats=np.asarray([[1.0, 1000.0, v] for v in range(4)], np.float32), + box=np.asarray([[0.0, -1.0, 0.0]], np.float32), + box_feats=np.asarray([[0.0, 1.0, 0.0]], np.float32), + ) + return dict(config=str(cfg), checkpoint=str(ckpt), inputs=str(data)) + + +def test_assets_preserve_training_feature_normalization(assets): + cfg, (features, box, normals), rotation = load_assets(assets) + values = np.arange(4, dtype=np.float32) + np.testing.assert_array_equal( + features[:, 0], (values - values.mean()) / (values.std() + 1e-8) + ) + np.testing.assert_array_equal(rotation, np.eye(3)) + assert cfg["resolved_other_feats_channels"] == 1 + assert box.shape == normals.shape == (1, 3) + + +@pytest.mark.parametrize( + "change,match", + [ + ( + {"model_to_world": [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]]}, + "proper orthonormal", + ), + ({"checkpoint": "/nonexistent/worldparticle.ckpt"}, "existing local file"), + ({"unexpected_option": 1}, "Unknown WorldParticle options"), + ], +) +def test_invalid_assets_fail_before_loading_network(assets, change, match): + with pytest.raises(ValueError, match=match): + load_assets(dict(assets, **change)) + + +def test_unsupported_checkpoint_task_is_rejected(assets): + Path(assets["config"]).write_text("data: unsupported_task\ntimestep: 0.02\n") + with pytest.raises(ValueError, match="supported data"): + load_assets(assets) + + +class _TranslationNetwork(torch.nn.Module): + + def forward(self, inputs, *, k_partite, external_force): + q, qd, *_ = inputs + next_qd = qd + 0.01 * external_force + return q + 0.005 * (qd + next_qd), next_qd + + +@pytest.mark.skipif( + not wp.is_cuda_available(), reason="Public fluid authoring requires CUDA" +) +def test_public_worldparticle_provider_advances_and_resets(assets, monkeypatch): + import dexsim + from dexsim.engine.newton_physics import NewtonCfg, ProviderNewtonSolverCfg + from embodichain.lab.sim.physics.worldparticle import WorldParticleProvider + from embodichain_models.dynamics.worldparticle import _runtime + from dexsim.spawn import SceneBuilder, FluidRenderMode + + monkeypatch.setattr( + _runtime, "load_network", lambda *_: _TranslationNetwork().cuda() + ) + config = dexsim.WorldConfig() + config.open_windows = False + config.init_ground = False + config.use_default_physics = False + config.backend = dexsim.types.Backend.VULKAN + config.renderer = dexsim.types.Renderer.NORENDER + rotation = [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + config.newton_cfg = NewtonCfg( + device="cuda:0", + gravity=0.0, + dt=0.01, + num_substeps=1, + sync_to_dexsim=False, + collision_pipeline_cfg=None, + solver_cfg=ProviderNewtonSolverCfg( + provider_factory=WorldParticleProvider, + options=dict(assets, model_to_world=rotation), + ), + ) + config.newton_cfg.use_cuda_graph = False + world = dexsim.World(config) + builder = SceneBuilder(world) + initial = np.asarray([[i, 0.0, 1.0] for i in range(4)], np.float32) + builder.add_fluid( + "fluid", + initial, + density=1000.0, + particle_spacing=0.1, + render_mode=FluidRenderMode.PARTICLES, + ) + scene = builder.finalize() + try: + fluid = scene.get_fluid_particle_set("fluid") + velocity = np.tile([0.0, 0.0, 2.0], (4, 1)) + fluid.set_particle_velocities(velocity) + world.update() + np.testing.assert_allclose( + fluid.get_particle_positions().numpy(), + initial + [0.0, 0.0, 0.02], + atol=1e-6, + ) + np.testing.assert_allclose( + fluid.get_particle_velocities().numpy(), velocity, atol=1e-6 + ) + assert scene.dynamics_capabilities().supported_particle_kinds == ("fluid",) + scene.reset() + np.testing.assert_allclose( + fluid.get_particle_positions().numpy(), initial, atol=1e-6 + ) + np.testing.assert_array_equal( + fluid.get_particle_velocities().numpy(), np.zeros((4, 3)) + ) + finally: + scene.close() + + +@pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA inference adapter") +def test_particle_force_units_rotation_buffers_and_invalid_dt(): + import newton + from dexsim.engine.newton_physics.particle_inference import ParticleInferenceSolver + from embodichain_models.dynamics.worldparticle.model import WorldParticleModel + + builder = newton.ModelBuilder() + builder.add_particle( + pos=wp.vec3(0.0, 0.0, 0.0), vel=wp.vec3(0.0, 0.0, 0.0), mass=2.0 + ) + model = builder.finalize(device="cuda:0") + model.set_gravity((0.0, 0.0, -10.0)) + first, second = model.state(), model.state() + wp.to_torch(first.particle_f).copy_( + torch.tensor([[4.0, 0.0, 0.0]], device="cuda:0") + ) + rotation = np.asarray( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], np.float32 + ) + solver = ParticleInferenceSolver( + model, + lambda device: WorldParticleModel( + _TranslationNetwork().cuda(), + (torch.zeros((1, 1), device="cuda:0"), None, None), + rotation, + {"timestep": 0.01, "k_partite": 2}, + ), + 0.01, + ) + try: + solver.step(first, second, model.control(), None, 0.01) + wp.synchronize_device(model.device) + np.testing.assert_allclose( + second.particle_q.numpy(), [[0.0001, 0.0, -0.0005]], atol=1e-7 + ) + np.testing.assert_allclose( + second.particle_qd.numpy(), [[0.02, 0.0, -0.1]], atol=1e-6 + ) + np.testing.assert_array_equal(first.particle_q.numpy(), np.zeros((1, 3))) + retained = second.particle_q.numpy().copy() + with pytest.raises(ValueError, match="step dt"): + solver.step(first, second, model.control(), None, 0.02) + np.testing.assert_array_equal(second.particle_q.numpy(), retained) + finally: + solver.close() + with pytest.raises(RuntimeError, match="closed"): + solver.step(first, second, model.control(), None, 0.01) + + +@pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA provider creation") +@pytest.mark.parametrize( + "change,match", + [ + ({"dt": 0.02}, "substep dt"), + ({"requires_grad": True}, "inference only"), + ], +) +def test_unsupported_runtime_rejected_before_weights( + assets, monkeypatch, change, match +): + import newton + from dexsim.engine.newton_physics import NewtonSolverContext + from embodichain.lab.sim.physics.worldparticle import WorldParticleProvider + from embodichain_models.dynamics.worldparticle import _runtime + + def forbidden_load(*_): + raise AssertionError("Weights allocated before runtime validation") + + monkeypatch.setattr(_runtime, "load_network", forbidden_load) + builder = newton.ModelBuilder() + for i in range(4): + builder.add_particle( + pos=wp.vec3(float(i), 0.0, 0.0), vel=wp.vec3(0.0, 0.0, 0.0), mass=1.0 + ) + model = builder.finalize( + device="cuda:0", requires_grad=change.get("requires_grad", False) + ) + model.set_gravity((0.0, 0.0, 0.0)) + context = NewtonSolverContext( + model=model, + control=model.control(), + contacts=None, + device=model.device, + dt=change.get("dt", 0.01), + scene_indices={}, + ) + with pytest.raises(ValueError, match=match): + WorldParticleProvider().create(context, assets) + + +@pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA inference adapter") +@pytest.mark.parametrize("mode", ["zero_mass", "inactive"]) +def test_runtime_fixed_particle_remains_fixed(mode): + import newton + from dexsim.engine.newton_physics.particle_inference import ParticleInferenceSolver + from embodichain_models.dynamics.worldparticle.model import WorldParticleModel + + builder = newton.ModelBuilder() + builder.add_particle( + pos=wp.vec3(0.0, 0.0, 1.0), vel=wp.vec3(0.0, 0.0, 0.0), mass=1.0 + ) + model = builder.finalize(device="cuda:0") + model.set_gravity((0.0, 0.0, -10.0)) + solver = ParticleInferenceSolver( + model, + lambda device: WorldParticleModel( + _TranslationNetwork().cuda(), + (torch.zeros((1, 1), device="cuda:0"), None, None), + np.eye(3, dtype=np.float32), + {"timestep": 0.01, "k_partite": 2}, + ), + 0.01, + ) + first, second = model.state(), model.state() + # Public particle inverse-mass writes can fix a particle after creation. + if mode == "zero_mass": + model.particle_inv_mass.zero_() + else: + model.particle_flags.zero_() + try: + solver.step(first, second, model.control(), None, 0.01) + wp.synchronize_device(model.device) + np.testing.assert_array_equal( + second.particle_q.numpy(), first.particle_q.numpy() + ) + np.testing.assert_array_equal(second.particle_qd.numpy(), np.zeros((1, 3))) + finally: + solver.close() + + +@pytest.mark.parametrize( + "task,raw,expected", + [ + ("fluid", [[0.125, 0.01], [0.25, 0.02]], [[0.125, 0.01], [0.25, 0.02]]), + ("duck_drop", [[0.2], [0.4]], [[0.2], [0.4]]), + ( + "cloth_on_ball", + [ + [1.0, 0.0, 0.0, 0.0, 0.0, 5000.0, 50.0], + [3.0, 0.0, 0.0, 0.0, 0.0, 300000.0, 500.0], + ], + [[-1.0, 0.0, 0.0], [1.0, 1.0, 1.0]], + ), + ], +) +def test_task_specific_feature_semantics(task, raw, expected): + from embodichain_models.dynamics.worldparticle._runtime import normalize_features + + np.testing.assert_allclose( + normalize_features(task, np.asarray(raw, np.float64)), expected, atol=1e-7 + ) + + +def test_topology_layout_and_capabilities_are_model_scoped(assets): + from embodichain.lab.sim.physics.worldparticle import WorldParticleProvider + + Path(assets["config"]).write_text( + "data: cloth_on_ball\ntimestep: .02\ngravity: [0.,0.,0.]\n" + ) + raw = np.tile([1.0, 0.0, 0.0, 0.0, 0.0, 5000.0, 50.0], (4, 1)) + values = dict( + feats=raw, + box=np.zeros((1, 3)), + box_feats=np.ones((1, 3)), + neighbors_index=np.asarray([1, 0, 3, 2]), + neighbors_row_splits=np.arange(5), + ) + np.savez(assets["inputs"], **values) + cfg, arrays, _ = load_assets(assets) + assert cfg["resolved_use_topology_conv"] is True + assert cfg["resolved_other_feats_channels"] == 3 + np.testing.assert_array_equal(arrays[3], values["neighbors_index"]) + cloth = WorldParticleProvider() + cloth.register_custom_attributes(None, assets) + assert cloth.capabilities.supported_particle_kinds == ("cloth",) + assert WorldParticleProvider().capabilities.supported_particle_kinds == ("fluid",) + values["neighbors_index"][0] = 4 + np.savez(assets["inputs"], **values) + with pytest.raises(ValueError, match="topology CSR"): + load_assets(assets) + + +def test_topology_checkpoint_requires_original_neighbors(assets): + Path(assets["config"]).write_text( + "data: duck_drop\ntimestep: .02\ngravity: [0.,0.,0.]\n" + ) + np.savez( + assets["inputs"], + feats=np.ones((4, 1)), + box=np.zeros((1, 3)), + box_feats=np.ones((1, 3)), + ) + with pytest.raises(ValueError, match="requires CSR"): + load_assets(assets) + + +def test_legacy_fluid_buffer_migration_does_not_hide_missing_weights(): + from embodichain_models.dynamics.worldparticle._runtime import load_checkpoint_state + + network = torch.nn.Module() + network.weight = torch.nn.Parameter(torch.ones(1)) + network.local_feature_extractor = torch.nn.Module() + network.local_feature_extractor.register_buffer( + "topology_filter_extent_tensor", torch.tensor(0.01) + ) + state = {"weight": torch.tensor([2.0])} + with pytest.warns(UserWarning, match="unused topology"): + load_checkpoint_state(network, state, legacy_fluid=True) + torch.testing.assert_close(network.weight, torch.tensor([2.0])) + torch.testing.assert_close( + network.local_feature_extractor.topology_filter_extent_tensor, + torch.tensor(0.01), + ) + with pytest.raises(RuntimeError, match="Missing key"): + load_checkpoint_state(network, {}, legacy_fluid=True) + with pytest.raises(RuntimeError, match="Missing key"): + load_checkpoint_state(network, state, legacy_fluid=False) + + +def test_explicit_provider_config_does_not_resolve_builtin_solvers(): + from embodichain.lab.sim.cfg import NewtonPhysicsCfg + from embodichain.lab.sim.physics.worldparticle import WorldParticleProvider + from dexsim.engine.newton_physics import ProviderNewtonSolverCfg + + provider = ProviderNewtonSolverCfg( + provider_factory=WorldParticleProvider, options={} + ) + cfg = NewtonPhysicsCfg( + solver_cfg=provider, collision_cfg=None, use_cuda_graph=False + ) + result = cfg.to_dexsim_cfg(gpu_id=0) + assert result.solver_cfg is cfg.solver_cfg + assert result.solver_cfg.provider_factory is WorldParticleProvider + assert result.collision_pipeline_cfg is None + assert result.use_cuda_graph is False diff --git a/tests/sim/objects/test_deformable_object.py b/tests/sim/objects/test_deformable_object.py index 0f0a0191e..dc516289e 100644 --- a/tests/sim/objects/test_deformable_object.py +++ b/tests/sim/objects/test_deformable_object.py @@ -183,12 +183,13 @@ def test_manager_rejects_non_particle_newton_solver(solver_type: str) -> None: sim.physics = NewtonPhysicsBackend(SimpleNamespace()) sim.physics._configured_solver_type = solver_type sim.device = torch.device("cuda") + sim.sim_config = SimpleNamespace(physics_cfg=NewtonPhysicsCfg()) with pytest.raises(NotImplementedError, match="does not support deformable"): sim.add_deformable_object(VolumeDeformableObjectCfg(uid="soft")) -@pytest.mark.parametrize("solver_type", ["auto", "dexuni"]) +@pytest.mark.parametrize("solver_type", ["auto", "dexuni", "provider", "worldparticle"]) @pytest.mark.parametrize( "config_type", [SurfaceDeformableObjectCfg, VolumeDeformableObjectCfg] ) @@ -199,12 +200,20 @@ def test_manager_declares_deformables_with_supported_solver( ) -> None: from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.spawn.scene import SpawnScene + from dexsim.engine.newton_physics import RegisteredNewtonSolverCfg sim = object.__new__(SimulationManager) sim.physics = NewtonPhysicsBackend(SimpleNamespace()) sim.physics._configured_solver_type = solver_type sim.device = torch.device("cuda") - sim.sim_config = SimpleNamespace(physics_cfg=NewtonPhysicsCfg()) + solver_cfg = ( + RegisteredNewtonSolverCfg(solver_type=solver_type) + if solver_type == "worldparticle" + else None + ) + sim.sim_config = SimpleNamespace( + physics_cfg=NewtonPhysicsCfg(solver_cfg=solver_cfg) + ) sim._deformable_objects = {} sim._visualization_topology_revision = 0 builder = SimpleNamespace( diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 0aee1b1b3..dd6be514f 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -1343,12 +1343,14 @@ def capture(self, **_kwargs: object) -> None: (False, True, "viser", False), ], ) +@pytest.mark.parametrize("enable_default_ground", [True, False]) def test_constructor_starts_visualization_after_default_scene( monkeypatch: pytest.MonkeyPatch, headless: bool, entity_gizmo: bool, backend: str, expected_gizmo: bool, + enable_default_ground: bool, ) -> None: lifecycle: list[str] = [] spawn_scene = MagicMock() @@ -1430,6 +1432,7 @@ def start_visualization(sim: SimulationManager) -> None: SimulationManagerCfg( startup_summary="off", num_envs=3, + enable_default_ground=enable_default_ground, headless=headless, enable_entity_gizmo=entity_gizmo, visualization=VisualizationCfg(backend=backend), @@ -1443,7 +1446,7 @@ def start_visualization(sim: SimulationManager) -> None: "arenas", "resources", "background", - "plane", + *(["plane"] if enable_default_ground else []), "lighting", "visualization:3", ] + (["entity_gizmo"] if expected_gizmo else [])