Make deformables Newton-only and limit differentiation to kinematics

agent_context/MAP.yaml

@@ -26,10 +26,29 @@ topics:
       - arena
       - physics step
       - sim update
+      - runtime control
+      - kinematic joint trajectory
+      - register_kinematic_joint_trajectory
+      - kinematic nodal trajectory
+      - register_kinematic_nodal_trajectory
+      - KinematicJointTrajectoryControl
+      - contact material schedule
+      - register_contact_material_schedule
+      - particle contact material schedule
+      - register_particle_contact_material_schedule
       - scene object
       - asset registry
       - manual update
       - GPU physics
+      - deformable
+      - particle set
+      - Newton soft body
+      - Newton cloth
+      - array-backed mesh
+      - stable node indices
+      - tetrahedral particle order
+      - independent visual mesh
+      - visual binding mode
       - simulation lifecycle
       - collision_policy
       - collision isolation
@@ -53,10 +72,12 @@ topics:
     source_of_truth:
       - embodichain/lab/sim/__init__.py
       - embodichain/lab/sim/sim_manager.py
+      - embodichain/lab/sim/_runtime_controls.py
       - embodichain/lab/sim/cfg/
       - embodichain/lab/sim/_legacy_cfg.py
       - embodichain/lab/sim/common.py
       - embodichain/lab/sim/material.py
+      - embodichain/lab/sim/shapes.py
       - embodichain/lab/sim/profiler.py
       - embodichain/lab/sim/spawn/descriptors.py
       - embodichain/lab/sim/spawn/scene.py
@@ -361,6 +382,9 @@ topics:
       - deformable
       - soft body
       - cloth
+      - Newton deformable
+      - particle set
+      - render topology
       - quaternion
       - xyzw
     paths:
@@ -594,6 +618,7 @@ topics:
     title: Differentiable Environment (APG)
     aliases:
       - differentiable env
+      - DifferentiableEnv
       - apg
       - analytic policy gradient
       - differentiable rl
@@ -608,14 +633,15 @@ topics:
       - warp tape
       - requires_grad
       - semi_implicit
-      - DifferentiableEmbodiedEnv
+      - DifferentiableEnv
       - NewtonStepFunc
       - quaternion
       - xyzw
     paths:
       - topics/differentiable-env/differentiable-env.md
     source_of_truth:
       - embodichain/lab/gym/envs/differentiable_env.py
+      - embodichain/lab/sim/cfg/simulation.py
       - embodichain/lab/sim/diff/
       - embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py
     related_topics:

agent_context/topics/differentiable-env/differentiable-env.md

@@ -1,114 +1,95 @@
 # differentiable-env
 
-> Topic: Differentiable environment for analytic policy gradient (APG) —
-> `DifferentiableEmbodiedEnv` + the `embodichain.lab.sim.diff` Warp-tape
-> ↔ PyTorch-autograd bridge.
+> Topic: Newton-backed kinematic environments for analytic policy gradient
+> (APG) through the Warp-tape ↔ PyTorch-autograd bridge.
+
+## Public entry point
+
+Use
+`embodichain.lab.gym.envs.differentiable_env.DifferentiableEnv`.
+It inherits `EmbodiedEnv`, preserves its scene lifecycle, and replaces the
+normal physics step with a task-defined kinematics callback recorded on a
+Warp tape.
+
+The resolution path is:
+
+    DifferentiableEnv.step(action)
+      → NewtonStepFunc.apply(action, sim_state)
+      → _apply_action_kernel(action_wp, tape)
+      → _make_kinematic_step_fn()()
+      → _read_outputs(final_state)
+      → Warp tape backward → action.grad
+
+## Invariants
+
+- The configured physics backend must be `NewtonPhysicsCfg`.
+- `NewtonPhysicsCfg.requires_grad` must be `True`.
+- Default physics and other backends fail during `DifferentiableEnv`
+  construction.
+- `DifferentiableEnv` always supplies a named kinematics callback; its bridge
+  contract has no dynamics mode, solver substeps, or control buffer.
+- The environment never invokes the configured Newton solver or collision
+  pipeline.
+- Newton gradient configuration still selects the semi-implicit solver, but
+  `DifferentiableEnv` never advances it.
+- Gradient mode disables Newton CUDA graph capture.
 
-## Overview
+## Subclass contract
 
-EmbodiChain supports analytic policy gradient (APG) via
-`embodichain.lab.gym.envs.differentiable_env.DifferentiableEmbodiedEnv`.
-The bridge wraps a Warp tape around one EmbodiChain physics step and
-exposes a `torch.autograd.Function`
-(`embodichain.lab.sim.diff.NewtonStepFunc`) so PyTorch-side `action`
-tensors get a gradient from `tape.backward()`.
+Task authors implement three hooks:
 
-## Required configuration
+- `_apply_action_kernel(action_wp, tape)` launches Warp work that maps the
+  PyTorch action bridge array into task-owned kinematic state.
+- `_make_kinematic_step_fn()` returns a zero-argument callback such as
+  `newton.eval_fk(...)`. The callback returns the state consumed by the output
+  hook.
+- `_read_outputs(final_state)` returns `obs`, `reward`, `terminated`, and
+  `truncated`, plus `_order` and `_grad_track` metadata used by
+  `NewtonStepFunc`.
 
-- `NewtonPhysicsCfg(requires_grad=True, solver_cfg={"solver_type": "semi_implicit"})`
-- `use_cuda_graph=False` (forced by dexsim when grad mode is on)
+There is no public `_apply_dynamics_action_kernel` or
+`differentiable_step_mode` extension point. Differentiable dynamics are a
+future feature, not a `DifferentiableEnv` capability.
 
-The default backend and any other Newton solver are rejected at
-construction time by `DifferentiableEmbodiedEnv._validate_diff_cfg`.
+## Autograd and reset rules
 
-Newton/Warp `body_q` transforms contain position followed by a native `xyzw`
-quaternion. This already matches EmbodiChain's quaternion convention, so the
-differentiable bridge and FK reward path must not reorder those four
-components. The Franka target pose likewise uses `xyz + xyzw`, with identity
-orientation `(0, 0, 0, 1)`.
+Action, kinematics, and gradient-producing output kernels must execute while
+the Warp tape is open. Each tracked output names its backing Warp array in
+`_grad_track`; an output mapped to `None` does not seed Warp backward.
 
-## Subclass contract
+A grad-tracked terminal step returns the terminal observation and exposes
+`requires_reset_after_backward` plus `deferred_reset_ids` in `info`. Reset
+those rows only after backward. A no-grad terminal step resets them
+synchronously.
+
+## Franka reference task
+
+`embodichain_tasks.special.franka_reach_apg.FrankaReachApgEnv` is the canonical
+example. Its path is:
+
+    action → new_joint_q → newton.eval_fk → body_q → reward kernel → action.grad
 
-Task authors implement two methods on `DifferentiableEmbodiedEnv`:
-
-- `_apply_action_kernel(action_wp, tape)` — launch a Warp kernel that
-  writes joint/body targets into `nm._control` while the tape is open.
-  The `action_wp` argument is a `wp.array(dtype=wp.float32,
-  requires_grad=True)` of shape `[num_envs * action_dim]`.
-- `_read_outputs(final_state)` — build the `obs` / `reward` /
-  `terminated` / `truncated` outputs as torch tensors via `wp.to_torch`
-  so the tape can record the dependency. Must return a dict with
-  `_order` (tuple of output keys) and `_grad_track` (mapping from output
-  key to the Warp array that backs its gradient, or `None` for outputs
-  that don't need grad).
-
-Optionally override `_make_step_fn()` to swap the per-substep advance
-function. The default uses `dexsim.engine.newton_physics.DifferentiableStepper.step`;
-the Franka APG example overrides it to call `newton.eval_fk` directly
-(see "FK bypass" below).
-
-See `embodichain_tasks.special.franka_reach_apg` for
-the canonical example.
-
-## Why reward must be computed inside the tape
-
-`NewtonStepFunc.forward` keeps the `wp.Tape` open while
-`obs_reward_fn(final_state)` runs. Reward must be computed by a Warp
-kernel that writes into a `wp.zeros(..., requires_grad=True)` array
-inside the tape; `wp.to_torch(reward_wp)` then returns a torch tensor
-that carries the tape's gradient. Computing reward in pure torch *after*
-the tape closes would detach it from the autograd graph and
-`action.grad` would come back as `None`.
-
-The same rule applies to any observation that needs to be
-grad-tracked: build it from `wp.to_torch` of a tape-tracked Warp array.
-
-## FK bypass for the Franka task
-
-The `semi_implicit` Newton solver does not propagate gradient through
-`joint_target_pos` to `body_q` (verified empirically; the reference
-implementation at `/root/sources/analytic_policy_gradients/envs/franka_reach_env.py`
-hits the same limitation and uses the same workaround). The Franka APG
-example overrides `_make_step_fn()` to call `newton.eval_fk(model,
-new_joint_q, joint_qd, fk_state)` directly, bypassing the dynamics
-solver. The grad path is then:
-
-    action → new_joint_q (action kernel) → eval_fk → body_q → reward kernel → reward_wp → tape.backward → action.grad
-
-The default `_make_step_fn` still uses the differentiable stepper, so
-envs whose reward depends on dynamics (not just FK) can use it — but
-they should verify the solver actually propagates grad for their
-control inputs before relying on it.
-
-## Functor autograd compatibility
-
-Reward/observation functors that compose torch operations on tensors
-obtained via `wp.to_torch` are automatically autograd-compatible.
-Functors that detour through CPU / NumPy break the graph; those need
-torch-only reimplementations for the differentiable path. For now, the
-differentiable env computes reward via a dedicated Warp kernel rather
-than reusing the standard reward-manager functors — a future task can
-audit and port functors as needed.
-
-## Memory
-
-Each step records `sim_steps_per_control` substeps into the tape. For
-long horizons or large `num_envs`, pass `truncate_backward_at=K` on the
-env config to split the tape and detach at chunk boundaries.
+The task snapshots live joint positions before opening the tape and writes the
+detached next joint state back after the bridge returns. It does not exercise
+Newton dynamics.
+
+## Dynamics boundary
+
+The public differentiable package exposes no solver stepper, trajectory, or
+gradient-rollout API. Add those capabilities as a separate future design when
+Newton dynamics are ready for end-to-end validation.
 
 ## Source of truth
 
-- `embodichain/lab/gym/envs/differentiable_env.py` —
-  `DifferentiableEmbodiedEnv` base class.
-- `embodichain/lab/sim/diff/bridge.py` — `NewtonStepFunc`,
-  `tape_context`, `differentiable_step`.
-- `embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py` —
-  example task.
-- `embodichain/lab/sim/sim_manager.py` —
-  `SimulationManager.create_differentiable_stepper` /
-  `create_gradient_rollout` delegators.
-- `/root/sources/dexsim/python/dexsim/engine/newton_physics/differentiable_stepper.py`
-  — the underlying dexsim primitive.
+- `embodichain/lab/gym/envs/differentiable_env.py`
+- `embodichain/lab/sim/cfg/simulation.py`
+- `embodichain/lab/sim/diff/`
+- `embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py`
+
+## Focused validation
+
+- `tests/gym/envs/test_differentiable_embodied_env.py`
+- `tests/sim/test_sim_manager_cfg.py`
 
 ## Related topics
 

agent_context/topics/sim-visualization/sim-visualization.md

@@ -206,8 +206,8 @@ Deformable vertices are stored relative to the corresponding arena node.
 | `RigidObjectGroup` | One node and pose per constituent object |
 | `Robot` | One mesh node per non-empty link |
 | `Articulation` | One mesh node per non-empty link |
-| Volume `DeformableObject` (`SoftObject`) | Live collision vertices with a cached convex-hull surface |
-| Surface `DeformableObject` (`ClothObject`) | Live physical vertices with render triangles mapped onto the welded physical vertex buffer |
+| Volume `DeformableObject` (`SoftObject`) | Live Newton render-surface vertices and triangles |
+| Surface `DeformableObject` (`ClothObject`) | Live Newton render-surface vertices and triangles |
 | `Camera` | Frustum plus optional low-frequency RGB preview |
 | Default ground | 1000 m × 1000 m XY grid, 1 m cells, 10 m sections |
 | `SceneOverlays` | Frames, targets, trajectories, and point clouds |
@@ -235,21 +235,28 @@ slow rendering or clients cannot accumulate an image backlog.
 
 ## Deformables
 
-Volume and surface deformables currently require Default-backend GPU physics.
-Their live vertices are sampled at `soft_body_fps`, independently from
-`scene_fps`. `SceneExporter` enumerates the manager's single deformable
-registry and reads both topologies through `get_surface_vertices()` and
-`get_surface_triangles()`; it does not branch on legacy buffer APIs. The
-`deformable_type` discriminator only selects the existing soft/cloth browser
-node kind, path, and color.
-
-- DexSim does not expose soft-body collision triangle connectivity.
-  `VolumeDeformableData.collision_surface_triangles` therefore caches a SciPy
-  `ConvexHull` over rest collision vertices. The preview follows deformation
-  but cannot preserve concave render detail.
-- Cloth maps all render-mesh triangles onto DexSim's welded rest-vertex buffer
-  with `cKDTree`. Construction raises `RuntimeError` if the mapping distance
-  exceeds the scale-relative tolerance.
+Volume and surface deformables require the Newton backend, CUDA, and a
+particle-capable solver. Their live vertices are sampled at `soft_body_fps`,
+independently from `scene_fps`. `SceneExporter` enumerates the manager's single
+deformable registry and reads render topology through
+`get_surface_vertices()` and `get_surface_triangles()`; it does not branch on
+legacy buffer APIs. The facade returns world-frame render vertices, and the
+exporter subtracts the arena offset before publishing them below the arena
+node. The `deformable_type` discriminator only selects the existing
+soft/cloth browser node kind, path, and color.
+
+- Both soft bodies and cloth publish the live render surface exposed by their
+  DexSim 0.5 typed Newton particle-set handles. No convex-hull reconstruction
+  or nearest-neighbor welding is performed in the visualization path.
+- A volume deformable separately exposes its tetrahedral collision surface
+  through `get_collision_surface_triangles()` for consumers that need physical
+  rather than render topology. `SceneExporter` intentionally uses render
+  topology.
+- Spawn binding validates that every replicated instance has the same render
+  vertex and triangle counts. A file-backed soft body whose DexSim clones have
+  a render topology different from the source fails during scene preparation;
+  use one environment or a compatible mesh until DexSim replication preserves
+  the source topology.
 - Viser does not update mesh vertices in place. `ViserBackend` removes and
   recreates a deformable mesh handle only when a dynamic vertex sample arrives.
   Pose-only frames reuse the current handle.
@@ -318,8 +325,7 @@ payload bytes plus capture/upload time.
 | Startup timeout or address-in-use error | The Viser worker did not become ready or the configured port is occupied. Select another port and inspect `visualization_health.worker_error`. |
 | Asset added after startup is missing | Step once, call `refresh_visualization()`, or mark topology dirty if the change bypassed manager APIs. |
 | Browser stops updating after an exporter/backend exception | `capture_visualization_safely()` latches the first error to protect simulation. Inspect health/logs, then stop and restart after fixing the cause. |
-| Soft body looks inflated or loses cavities | The surface is a collision-vertex convex hull, not the render topology. |
-| Cloth construction raises a mapping error | Render vertices do not match the welded physical rest vertices within tolerance. |
+| Replicated soft body fails with a render-vertex-count mismatch | DexSim produced clone render topology different from the source. Use one environment or a compatible mesh until DexSim replication preserves the topology. |
 | Camera frustum exists but preview is blank | Color capture is disabled, no image has been captured yet, or the selected camera/environment is hidden. |
 | Stereo/contact sensor is absent | Current camera export accepts only `sensor_type == "Camera"`; non-mesh sensors are not exported. |
 | Browser lags or upload cost is high | Reduce scene/image/deformable FPS, select fewer environments, or lower point-cloud limits. |