cd /news/robotics/real-time-inference-for-jepa-world-m… · home › topics › robotics › article
[ARTICLE · art-168447] src=anyscale.com ↗ pub= topic=robotics verified=true sentiment=↑ positive

Real-time inference for JEPA world model architectures with Ray

A design using a Ray cluster cuts V-JEPA2-AC planning latency from 16 seconds per action on a single on-board RTX 4090 to under one second, according to a November 12, 2026 post by Omar Shorbaji and Ian Jordan, PhD. The V-JEPA2-AC world model, which reports 65–80% zero-shot pick-and-place success on Franka arms with novel objects, spends roughly 8,000 predictor evaluations per action across ten CEM rounds of 800 candidate action sequences, and the reference controller blocks motion while planning. The authors move world model inference to NVIDIA A100-40GB GPUs in the cloud and distribute the search across a Ray cluster, noting that a single optimized A100 alone measured 21.4 seconds per action.

by read14 min views1 publishedNov 12, 2026
Real-time inference for JEPA world model architectures with Ray
Image: Anyscale (auto-discovered)

Omar Shorbajiand

Ian Jordan, PhD| November 12, 2026

With V-JEPA2-AC the system decides what a robot arm should do next by assessing hundreds of possible futures and picking the one closest to the goal. It takes a significant amount of time between moves. This document presents a design and tested implementation for getting that under a second by moving the search into a Ray cluster.

Link01 · A Primer on Robotics #

A robot perceives a partial observation of the environment, selects an action via some policy, executes it through a controller, and repeats this at control-loop frequency. Classically, each of those three functions is solved in closed form, and a CPU on the robot is enough to run all of them.

Classical robotics built each of those three boxes from first principles: hand-derived kinematics for perception, trajectory optimization for planning, PID or model-predictive controllers for control. This approach, depicted in Figure 1, underlies decades of industrial automation, but every new task, object, or environment typically requires new engineering.

Classical robotics — ~1 ms / action

Physical AI aims to replace some or all of those boxes with networks trained on data, trading hand-derived precision for generality: a model trained on sufficient data transfers to situations not explicitly programmed for.

LinkHow the JEPA architecture fills the "plan" box

V-JEPA2-AC is one such model. It is built on the Joint-Embedding Predictive Architecture (JEPA). In this architecture, the system decides what a robot arm should do next by assessing hundreds of possible futures using the model and picking the one closest to the goal. Figure 2 shows the whole loop in place on the hardware that runs it.

Deployed zero-shot on Franka arms in two labs, with no fine-tuning for the specific environment, V-JEPA2-AC reports 65–80% success on pick-and-place with novel objects. This ran on a single RTX 4090. However, that configuration takes 16 seconds to plan every single action. It is not fast enough to run continuously.

LinkWhere the sixteen seconds actually goes

It is worth walking one planning step end to end, because the arithmetic explains the rest of this document. The arm holds still. The camera frame and the goal image are encoded once. The search algorithm, CEM, then draws 800 candidate action sequences from a Gaussian centered on its current best guess, pushes every one through the predictor, scores the results against the goal embedding, keeps the best 80, and re-centers the Gaussian on them. That is one round, and it happens ten times: roughly 8,000 predictor evaluations for one small movement of one arm. The system then executes only the first action of the winning sequence, throws the rest away, and starts over from a fresh frame. The cost is structural, and it recurs at every step.

V-JEPA2-AC on one on-board GPU — 16 s / action

The controller is also blocking: it waits for the previous action to complete before sending the next, so planning never overlaps with motion. The robot s, plans for sixteen seconds, moves, and s again. That stutter, not the model's accuracy, is the problem this document sets out to fix.

Link02 · The Approach #

We take the approach of, first, moving world model inference off the robot and into a cloud server. While this introduces the new cost of a network round trip on every control step, that is offset by other improvements. First, it opens the opportunity of serving the model on a more powerful GPU. In this implementation we use NVIDIA A100-40GB GPUs, which have more memory than the 24 GB RTX 4090 used in the paper and substantially more bf16 throughput.

The same search on one cloud GPU — 21.4 s / action (measured, A100, optimized model)

On its own, a bigger GPU is not enough. Meta's reference planner, run unmodified on one A100 with these settings, takes 81.8 s per action; with the model-level optimizations in Section 05 it takes 21.4 s. Our 81.8 s is slower than the paper's 16 s on a 4090 because we ran Meta's released planning code as-is (fp32, 800 candidates × 10 rounds, a two-step horizon); the paper's timing likely reflects different settings or precision. Every speedup in this post is measured against our own 81.8 s baseline on the same hardware. Faster, but the arm still stops between moves.

The second step is to distribute the load across multiple GPUs. The opportunity comes from noting that within a single CEM refinement round rollouts can be computed independently and gathered before the round moves on. The paper runs them as one large batch on one GPU; nothing requires them to stay on one GPU. In our implementation we use one eight-GPU node (8× A100-40GB) on Anyscale.

The search distributed across a Ray cluster — 2.86 s / action (measured, 8× A100); 0.86 s with 3 rounds per step

We implement this using Ray.

Link03 · A Primer on Ray #

LinkA general framework for turning Python into a cluster program

Ray is an open-source framework, originally built at UC Berkeley's RISELab and now maintained by Anyscale, for scaling ordinary Python from a laptop to a cluster without rewriting it as a distributed system. It exposes a small set of primitives and preserves mostly-sequential code structure.

LinkRay Core

Ray Core has two building blocks that matter here. A @ray.remote function becomes a task, scheduled somewhere in the cluster and returning a future. A @ray.remote class becomes an actor: a stateful worker pinned to a machine, and optionally a GPU, that keeps loaded model weights resident between calls. Tasks suit one-off parallel work; actors suit a process that pays to load a 1.2-billion-parameter checkpoint once and then answers many cheap requests.

LinkRay Serve

Ray Serve sits on top of Ray Core: a @serve.deployment declaration is enough for Ray to manage a pool of replicas, route requests, and scale under load. Ray Core answers how to run a function on many workers; Ray Serve answers how to expose it as a service other systems call, which is the shape a robot's control loop needs, a low-latency endpoint hit once per step.

Link04 · Implementation #

LinkTwo Ray Serve deployments, composed

The design is two Ray Serve deployments and a handle between them. WorldModel is the GPU tier: one replica per GPU, each holding the encoder and predictor (we ran 8 replicas on 8 A100s). Planner is the driver: it runs the ten-round CEM loop, fans each round out across those replicas, and serves the endpoint the robot calls once per control step. Composing them this way, rather than as loose Ray Core actors, puts them under Serve's control: placement, health checks, restarts, and a declarative config a KubeRay RayService can roll out without downtime. The shape below elides the internals; what matters is where the fan-out sits. See figure 5 for a diagram of the deployment layout.

serve_app.py — Ray Serve · the shape, details elided

from ray import serve
Import ray
from ray.serve.handle import DeploymentHandle

@serve.deployment(
    num_replicas=8,  # one per GPU
    ray_actor_options={"num_gpus": 1},
    max_ongoing_requests=1,          # one shard at a time, so fan-out maps 1:1
)
class WorldModel:
    def __init__(self, checkpoint): ...        # paid once, when Serve starts it
    def encode(self, frame, proprio): ...      # the encoder is part of the model
    def score(self, ctx, actions): ...   # one shard, one GPU batch

@serve.deployment(num_replicas=2)
@serve.ingress(app)
class Planner:
    def __init__(self, model: DeploymentHandle):
        self.model = model

    @app.post("/plan")
    async def plan(self, frame, proprio, goal_image):
        obs, goal = await asyncio.gather(      # encode once per control step
            self.model.encode.remote(frame, proprio),
            self.model.encode.remote(goal_image),
        )
  ctx = ray.put((obs, goal))  # once per plan, read from shared memory
        mean, std = init_distribution(horizon, action_dim)

        for _ in range(10):                    # CEM rounds: sequential, the floor
            cands = sample(mean, std, n=800)
            scores = await asyncio.gather(*[   # fan out one shard per replica
                self.model.score.remote(ctx, shard)
                for shard in np.array_split(cands, 8)
            ])                                 # gather before the round refits
            mean, std = refit(cands, scores, elite_frac=0.1)

        return {"action": mean[0].tolist()}    # first action only, then re-plan

app_builder = Planner.bind(WorldModel.bind(CHECKPOINT))

Two choices carry the design. The GPU tier is pinned at num_replicas=8 rather than autoscaled, because a search whose replica count changes mid-round has unpredictable latency, and the point is a stable per-step budget. max_ongoing_requests=1 makes each replica take one shard at a time, so an 8-way fan-out lands one shard per GPU instead of queueing several on a subset. Underneath, each replica is still a Ray Core actor, which is why the weights stay resident; the difference is that Serve owns the lifecycle.

What we learned building it. Three details decided the latency. First, the observation and goal latents (2.8 MB) go into Ray's object store once per plan with ray.put; every replica reads them from shared memory and caches them on its GPU, so each round moves only about 5 KB of candidate actions out and the energies back, and Ray's whole share of a round (fan-out, gather, refit) is under 10 ms. Second, with max_ongoing_requests=1 Serve's default power-of-two-choices router occasionally placed two shards on one replica while another sat idle, doubling that round; Serve's round-robin request router lands the shards of a round on consecutive replicas, one per GPU (zero collisions in 50 rounds). Third, our Serve replica processes started with TORCHDYNAMO_DISABLE=1 in their environment, which silently turned torch.compile into a no-op; clearing it in the replica's constructor recovered a 1.35× speedup.

Source. The full implementation, including model , the fast scoring path, the complete CEM refinement step, the two Serve deployments, every benchmark script and a runnable notebook, lives in the companion repository: https://github.com/anyscale/robotics_world_model_inference/tree/main/v3_realtime_inference_for_jepa

Link05 · Results #

LinkMeasured on 8× A100

We built the design above and measured it on one Anyscale node with eight A100-40GB GPUs, using the search settings from Section 01 (800 candidates × 10 rounds, keep 80, horizon 2). The anchor is Meta's reference planner, run unmodified in fp32 on one A100: 81.8 s per action. Before adding GPUs we made the model itself faster on one GPU: a KV cache across rollout steps (the predictor is frame-causal, so earlier frames never need recomputing), bf16 with FlashAttention, and torch.compile. The result matches the reference energies (Pearson r = 0.9999) and brings one GPU to 21.4 s per action, 4× faster.

Then the fan-out. Each round's 800 candidates are split across W WorldModel replicas, one GPU each. Going from 1 to 8 GPUs cuts a plan from 21.4 s to 2.86 s, a 7.5× speedup. Same seeds give identical plans at every GPU count, so sharding changes the speed of the search but not its answer. The latency model proposed in the first draft of this post holds up: calibrated on our measured single-GPU round (2.14 s), it predicts 2.94 s at 8 GPUs against 2.86 s measured.

LinkPlanning latency vs. cluster size

A100-40GB replicas on one 8-GPU node. Points are measured; the line is the latency model.

2.86 s / action at 8 A100 workers. This is 7.5× faster than one A100 and 29× faster than Meta's reference planner.

Model: time(W) = 10 × ( r / effective(W) + 0.001s × (W − 1) ), where effective(W) = W / (1 + 0.008×(W−1)) and r is the measured single-GPU round, 2.14 s on an A100. Beyond 8 GPUs the curve flattens, because each GPU's share of the 800 candidates gets small. Larger clusters pay off mainly by serving more robots, or larger searches, from the same tier.

Two things remain unmeasured: network latency between an on-premises robot and a cloud region, and performance on a physical arm (our closed-loop tests use the world model as the simulator). Cold start stays off the request path: replicas load and compile in about 50 s at deploy time, before they take traffic. Batching candidates differently does not change accuracy: the plans are identical at 1, 2, 4 and 8 GPUs.

Link06 · Further Optimization #

LinkWhat to try once the search itself is no longer the bottleneck

Once each round is as fast as the GPUs allow, what remains is the chain of ten sequential rounds, and the next gains come from restructuring the loop rather than adding GPUs.

LinkA. Non-blocking control

Start planning step t+1 while the robot is still executing step t. That hides planning latency behind execution time instead of adding to it, and it is a controller change, not a model change.

Measured. We simulated an arm that takes 1 s per action and lands with 4 mm of execution error. A blocking controller spent 1.86 s per step (plan, then move). Non-blocking control, planning step t+1 from the world model's prediction of where step t will land, spent 1.04 s per step, i.e. the motion itself plus 40 ms. It reached the same final accuracy, and a 12-step episode took 15.3 s instead of 25.3 s.

LinkB. Warm-started CEM

Receding-horizon control throws away everything but the first action. Re-using the rest to seed the next step's distribution, instead of a fresh Gaussian, lets each replan converge in fewer than ten rounds, shrinking the sequential chain parallelism cannot touch.

Measured. In closed-loop tests, 3 rounds per step (0.86 s on 8 A100s) reached the same accuracy as the full 10 (2.83 s): both settle about 4 cm from a goal that starts 21.6 cm away. Warm-starting from the previous plan worked, but on this task a fresh 3-round search did equally well, so the gain came from running fewer rounds. How many rounds a task needs is worth tuning per task.

LinkC. Fleet-level batching

A GPU tier sized for one robot is idle between its requests. Behind a Serve endpoint the same tier serves a fleet, and request batching can pack shards from different robots into one GPU batch, amortizing the same warm GPUs across many arms.

LinkD. Model compression

Distilling or quantizing the predictor shrinks the compute term itself, independent of how many workers it is split across. It stacks with distribution rather than competing with it.

LinkE. Serving-level optimization

The predictor rolls forward step by step over the horizon, so what LLM serving already relies on carries over: KV caching across rollout steps, prefill and decode disaggregation, continuous batching. These shrink the same compute term from a different direction, inside the same Serve deployment.

Implemented. KV caching across rollout steps, together with bf16, FlashAttention and torch.compile, is part of the measured system. It gives 4× on a single GPU and matches the reference energies to r = 0.9999.

Link07 · Takeaways #

The sixteen seconds is not an implementation defect. It follows from a deliberate architectural choice: a JEPA model predicts state rather than actions, so an action has to be searched for at inference time instead of read off a policy head. That is the same trade agentic systems make when they sample many candidates and score them, and it comes with the same benefit, a system that can be pointed at a new goal without retraining.

The four setups make the trade concrete. A hand-written policy buys millisecond actions and pays for them in task-specific engineering. V-JEPA2-AC trades that speed for a system that accepts any goal you can photograph. Moving the search to a cluster buys the speed back without giving up the generality, and does it in infrastructure rather than in the model. Measured on one node of eight A100s, Ray Serve takes the full search from 82 s to 2.9 s per action, with distribution doing most of the work and Ray adding under 10 ms per round. The ten sequential rounds set the floor that more hardware cannot move. Attacking that floor with fewer rounds per step brings it to 0.86 s, and non-blocking control hides what remains behind the arm's own motion.

LinkReferences

  1. Assran, M. et al. V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning. arXiv:2506.09985, 2025.arxiv.org/abs/2506.09985
  2. Meta AI. Introducing the V-JEPA 2 world model and new benchmarks for physical reasoning.ai.meta.com/blog/v-jepa-2-world-model-benchmarks
  3. Ray Project. Ray Core: Key Concepts, Tasks and Actors.docs.ray.io/en/latest/ray-core
  4. Ray Project. Ray Serve documentation.docs.ray.io/en/latest/serve
  5. Ray Project. KubeRay: RayService.docs.ray.io/en/latest/cluster/kubernetes

Planning-time and success-rate figures are as reported in the V-JEPA2-AC paper (arXiv:2506.09985). Cluster sizing, the CEM parallelization design, the latency model, and all code shown are an independent engineering implementation built on top of that published work, not affiliated with or endorsed by Meta AI or the Ray project. Latency measurements in Sections 05 and 06 were taken on one 8× A100-40GB node on Anyscale (November 2026); the extension of Figure 6 beyond 8 GPUs is a projection.

── more in #robotics 4 stories · sorted by recency
── more on @ray 3 stories trending now
sponsored brought to you by zahid.host 4,200+ EU-deployed projects
reading about agents? ship yours in a single git push.

Run your AI side-project on zahid.host

EU-based hosting, git-push deploys, automatic HTTPS, no cold starts. Free tier with a custom domain — perfect for shipping the agent you just read about.

$git push zahid main
→ Live at https://your-agent.zahid.host ✓
Get free account → Pricing
from €0/mo · no card required
LIVE [news/real-time-inference-…] indexed:0 read:14min 2026-11-12 · —