Fix JAX inference timing to include GPU sync

[Abhishek21g]

Summary

Fixes #983 — policy_timing.infer_ms under-reported GPU inference time because the timer stopped at JAX async dispatch instead of waiting for GPU work to complete.

Root cause

On the JAX path, Policy.infer() measured wall time around _sample_actions() but did not synchronize device work before stopping the timer. JAX returns from jitted calls once work is scheduled, not when kernels finish — reporters saw ~30ms gaps vs the JAX profiler.

PyTorch path is unaffected (.detach().cpu() already forces sync).

Fix

Call jax.block_until_ready(actions) on the JAX path before recording model_time.

Verification

uv run pytest src/openpi/policies/policy_test.py::test_jax_async_dispatch_under_reports_without_sync -q
uv run pytest src/openpi/policies/policy_test.py::test_jax_policy_infer_syncs_actions_before_timing -q

On GPU hardware with the policy server running:

uv run examples/simple_client/main.py --env ALOHA_SIM --num-steps 50
# policy_infer_ms should be within ~5ms of JAX profiler kernel time