# Real-time inference for JEPA world model architectures with Ray

> Source: <https://anyscale.com/blog/real-time-inference-for-jepa-world-models-ray-anyscale>
> Published: 2026-11-12 09:00:00+00:00

# Real-time inference for JEPA world model architectures with Ray

[Omar Shorbaji](https://anyscale.com/blog?author=omar-shorbaji)and

[Ian Jordan, PhD](https://anyscale.com/blog?author=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.

## Link**01 · 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.

### Link**How 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.

### Link**Where 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 pauses, plans for sixteen seconds, moves, and pauses again. That stutter, not the model's accuracy, is the problem this document sets out to fix.

## Link**02 · 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.

## Link**03 · A Primer on Ray**

### Link**A general framework for turning Python into a cluster program**

[__Ray__](https://docs.ray.io/en/latest/ray-core/walkthrough.html) 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.

### Link**Ray 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.

### Link**Ray 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.

## Link**04 · Implementation**

### Link**Two 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

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

# The GPU tier. One replica per GPU, pinned, weights resident between calls.
@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

# The driver. Holds a handle to the GPU tier and runs the search.
@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 loading, 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__

## Link**05 · Results**

### Link**Measured 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.

### Link**Planning 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.

## Link**06 · Further Optimization**

### Link**What 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.

#### Link**A. 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.

#### Link**B. 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.

#### Link**C. 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.

#### Link**D. 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.

#### Link**E. 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.

## Link**07 · 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.

### Link**References**

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__](https://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__](https://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__](https://docs.ray.io/en/latest/ray-core/walkthrough.html)
4. Ray Project. *Ray Serve documentation.*[__docs.ray.io/en/latest/serve__](https://docs.ray.io/en/latest/serve/index.html)
5. Ray Project. *KubeRay: RayService.*[__docs.ray.io/en/latest/cluster/kubernetes__](https://docs.ray.io/en/latest/cluster/kubernetes/index.html)

*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.*