Apex 0.4 Cross-Attention GFM
Apex 0.4 is a short-horizon bimanual robot motion model trained as a dynamic motor prior, not as a monolithic VLA policy. It combines a Q35 vision-language backbone with a LoRA adapter and a cross-attention Geodesic Flow Matching head that predicts smooth TCP motion for a dual-arm robot.
This is the best Apex 0.4 checkpoint from the current Giovanna run:
run: outputs/apex04_runs/apex_0_4_cross_gfm_full_mix_2000
best validation step: 1200
best val_flow_loss: 0.2283533364534378
records indexed: 461876
train records: 452639
validation records: 9237
action spec: apex_action_20d_bimanual_tcp_scalar_gripper_v1
The run reached step 1350 before stopping on a corrupted source JPEG in the old dataset mix. The best validation checkpoint was already saved before that failure and is the recommended checkpoint to use.
What The Model Does
Apex 0.4 receives:
3 images: overhead/top, left wrist, right wrist
instruction text
embodiment/context prompt
current left/right TCP pose and gripper state
It produces a 10-step short-horizon bimanual action trajectory:
left TCP delta translation
left TCP delta rotation as 6D rotation representation
right TCP delta translation
right TCP delta rotation as 6D rotation representation
left scalar gripper delta
right scalar gripper delta
The output shape is:
[10, 20]
Apex 0.4 predicts the action trajectory through a 14D flow velocity head:
left translation velocity: 3
left SO(3) tangent velocity: 3
right translation velocity: 3
right SO(3) tangent velocity: 3
left/right gripper velocity: 2
The flow head cross-attends to selected Q35 internal K/V activations instead of using only a pooled VLM context vector. In practice, this made Apex 0.4 stronger than the earlier Apex 0.3 pooled-context GFM run.
Files In This Repository
apex04_flow_head.pt
Best validation Apex 0.4 cross-attention GFM head checkpoint.
q35_lora_adapter/
PEFT LoRA adapter for the Q35 base model.
apex_0_4_manifest.json
Training configuration and model dimensions.
metrics.jsonl
Training and validation metrics from the run.
examples/infer_apex04_robot.py
Example Python inference script.
The Q35 base checkpoint is not included in this repository. The local training run used:
local_checkpoints/q35-fast-run02
To run inference, you need that compatible Q35 base checkpoint plus the LoRA adapter and Apex flow head from this repo.
Training Summary
Key configuration:
model_dim: 512
layers: 6
heads: 2
head_dim: 256
kv_heads: 2
lora_rank: 8
flow_lr: 3e-5
lora_lr: 5e-6
grad_accum: 16
dtype: float16
device: cuda
Best validation row:
{
"step": 1200,
"loss": 0.20997606217861176,
"val_flow_loss": 0.2283533364534378,
"val_samples": 128.0,
"loss_gfm_trans": 0.027616772829787806,
"loss_gfm_rot": 0.0446991182398051,
"loss_gfm_gripper": 0.019184564895113,
"loss_gfm_smooth": 2.3695120215415955
}
Install
From an Apex checkout:
git clone https://github.com/vitrus-ai/Apex.git
cd Apex
python -m venv .venv-apex
. .venv-apex/bin/activate
pip install torch transformers peft pillow huggingface_hub
You also need the same Q35-compatible base model used for training. On Giovanna this was:
local_checkpoints/q35-fast-run02
Quick Python Inference
The example script downloads the Apex 0.4 adapter and flow head from Hugging Face, loads a local Q35 base model, runs Q35 over three images plus a task prompt, samples a short-horizon GFM action trajectory, and writes the result as JSON.
python examples/infer_apex04_robot.py \
--repo-id lucas-vitrus/apex-0-4 \
--base-model local_checkpoints/q35-fast-run02 \
--overhead path/to/overhead.jpg \
--left-wrist path/to/left_wrist.jpg \
--right-wrist path/to/right_wrist.jpg \
--instruction "pick up the bottle and place it in the bin" \
--current-tcp-json current_tcp.json \
--output-json apex04_action.json
current_tcp.json should look like this:
{
"left": {
"x_m": 0.42,
"y_m": 0.18,
"z_m": 0.31,
"roll_rad": 0.0,
"pitch_rad": 1.57,
"yaw_rad": 0.0,
"gripper": 0.5
},
"right": {
"x_m": 0.42,
"y_m": -0.18,
"z_m": 0.31,
"roll_rad": 0.0,
"pitch_rad": 1.57,
"yaw_rad": 0.0,
"gripper": 0.5
}
}
The script writes a normalized action tensor and a more robot-friendly list of per-step deltas. Translation deltas are reported in millimeters because Apex 0.4 was trained with the Apex delta-window action contract.
Connecting A Robot
Apex 0.4 should sit above your robot's safety and control stack. It is not a low-level torque controller.
A practical runtime loop is:
1. Capture synchronized overhead, left wrist, and right wrist images.
2. Read the current left/right TCP pose and gripper aperture.
3. Send images + instruction + current TCP to Apex 0.4.
4. Receive a 10-step TCP delta trajectory.
5. Convert TCP deltas into your robot's control frame.
6. Run collision checks, workspace limits, speed limits, and IK.
7. Execute only the first one or few steps.
8. Re-observe and call Apex again in receding-horizon control.
Pseudo-code:
while task_is_active:
observation = robot.observe()
action = apex04.predict(
overhead=observation.overhead_image,
left_wrist=observation.left_wrist_image,
right_wrist=observation.right_wrist_image,
instruction=task_text,
current_tcp=observation.current_tcp,
)
safe_plan = safety_layer.filter_tcp_delta_plan(action["robot_steps"])
joint_plan = ik.solve_bimanual_tcp_plan(safe_plan)
robot.execute(joint_plan[:2])
Important integration notes:
- Convert coordinate frames before execution. Apex uses the canonical Apex TCP convention from the training data.
- Treat gripper output as scalar aperture delta, not dexterous finger torque.
- Use a safety layer for workspace bounds, self-collision, object collision, velocity, acceleration, and force limits.
- Use receding-horizon execution. Do not blindly execute all 10 predicted steps without re-observing.
- For a real Vitrus hand, use a separate hand/contact controller until Apex 0.5+ tokenized hand conditioning is trained and validated.
Limitations
- Apex 0.4 was trained on converted ABC/HY-style bimanual delta windows, not on direct real-time robot torque control.
- The model predicts short-horizon TCP and scalar gripper deltas only.
- It does not output joint angles, torques, force targets, or contact-rich dexterous commands.
- It expects a compatible Q35 base model and the Apex code used to instantiate the cross-attention GFM head.
- It must be wrapped in robot-specific frame conversion, IK, safety checks, and execution code.
Intended Use
Use Apex 0.4 as a research checkpoint for short-horizon bimanual TCP motion generation and as a baseline for Apex 0.5 structured field-token training.
Do not use it as an unsupervised autonomous robot controller in an unconstrained physical environment.