ckwolfe-research/hj-transformer-v4

HJ-Pretrained Affordance-Aware Transformer

Causal Transformer pretrained on Hamilton-Jacobi (HJ) reachability labels for multi-agent drone navigation. The model learns feasibility structure from HJ value functions: "is it feasible to close the gap?", "when should I commit to passing?", etc.

Model Architecture

Parameter	Value
d_model	128
n_heads	4
n_layers	4
d_ff	512
L (history steps)	16
K (max objects)	8
H (waypoint horizon)	10
dropout	0.1

Input Schema

Per-timestep tokens (L steps, 3+K tokens per step):

Token	Dim	Contents
Ego	9	x, y, z, vx, vy, vz, roll, pitch, yaw
Goal	7	Δpx, Δpy, Δpz, Δvx, Δvy, Δvz, goal_type
Map	16	Placeholder (zeros)
Object 1…K	8	Δpx, Δpy, Δpz, Δvx, Δvy, Δvz, radius, Δt_seen

Output Heads

Head	Shape	Description
WaypointHead	(H, 3)	Ego waypoints
AgentWaypointHead	(K, H, 3)	Per-agent centroid waypoints
RadiusHead	(K,)	Per-agent radius estimate

Training

Parameter	Value
Optimizer	AdamW
Learning rate	0.001
Weight decay	0.0001
LR schedule	OneCycleLR (cosine)
Batch size	64
Max epochs	10
Grad clip	1.0
λ_wp	1.0
λ_obj	0.5
λ_r	0.2

Data Generation

Training data is generated offline via the HJ-Warp pipeline:

1. HJ Solver: Lax-Friedrichs on 6D grid (Warp GPU kernels)
2. Domain Randomization: mass, drag, thrust, sensor noise, obstacle count
3. Teacher: HJ-conditioned safe waypoints (blend safe + goal directions)
4. Labels: safe_set, time_to_boundary, safe_direction, collision_within_H

Intended Use

Pretraining stage for multi-agent drone navigation. The pretrained model provides a warm start for downstream RL (PPO/SAC in Isaac Sim) or DAgger.

Limitations

• Model is pretrained only (no RL fine-tuning yet)
• HJ solver assumes pairwise interactions (no N-body coupling)
• Map token is a placeholder (no real occupancy grid)