Reduce the latency by increase the batch size for vision transformer

[xiaopoc]

Description of the change:
Currently, each image in obs.images in pi0.py is processed sequentially by the vision transformer in embed_prefix(), which leads to redundant kernel launches and increased runtime.

Motivation:
Batching all images together and encoding them in one forward pass reduces kernel launch overhead and enables better fusion. This can lead to lower latency during inference (~5ms speed up on RTX 4090).

Proposed Change:

• Stack images across cameras along a new axis, e.g., stacked_images = jnp.stack(list(obs.images.values()), axis=1)
• Flatten and encode all images in one forward pass via self.PaliGemma.img(...)
• Update the token aggregation and attention masks accordingly