diff --git a/assets/logo.png b/assets/logo.png
new file mode 100644
index 0000000..4e74a09
Binary files /dev/null and b/assets/logo.png differ
diff --git a/assets/pushtheobjtogoal.mp4 b/assets/pushtheobjtogoal.mp4
new file mode 100644
index 0000000..77188de
Binary files /dev/null and b/assets/pushtheobjtogoal.mp4 differ
diff --git a/docs/FLOWMATCHING.md b/docs/FLOWMATCHING.md
new file mode 100644
index 0000000..c928831
--- /dev/null
+++ b/docs/FLOWMATCHING.md
@@ -0,0 +1,5 @@
+# Flow Matching
+
+Flow Matching (FM) is a generative modeling paradigm that learns a continuous velocity field that transforms samples drawn from a simple source distribution (e.g., Gaussian noise) into samples from the target distribution (e.g., robot actions). Instead of learning to denoise noisy samples (as in diffusion models), flow matching directly regresses the velocity at each point in time that moves a sample along a probability path.
+
+ref. https://arxiv.org/abs/2210.02747
diff --git a/envs/metaworld_env.py b/envs/metaworld_env.py
index aef7bc1..b59acf5 100644
--- a/envs/metaworld_env.py
+++ b/envs/metaworld_env.py
@@ -17,10 +17,13 @@ def __init__(self, env_name='push-v3', seed=42, render_mode='rgb_array', camera_
             camera_name=camera_name
         )
         self.render_mode = render_mode
+        self.camera_name = camera_name
 
         obs, _ = self.env.reset()
         self.state_dim = self._extract_state(obs).shape[0]
         self.action_dim = self.env.action_space.shape[0]
+        self.action_low = np.asarray(self.env.action_space.low, dtype=np.float32)
+        self.action_high = np.asarray(self.env.action_space.high, dtype=np.float32)
         self.obs_shape = self._get_image().shape
 
     def _extract_state(self, obs):
@@ -49,11 +52,45 @@ def _extract_state(self, obs):
             state = obs
         return np.asarray(state, dtype=np.float32)
 
-    def _get_image(self):
-        img = self.env.render()
+    def _get_image(self, camera_name=None):
+        if camera_name is None or camera_name == self.camera_name:
+            img = self.env.render()
+        else:
+            try:
+                img = self.env.render(camera_name=camera_name)
+            except TypeError:
+                # Older wrappers may only support the camera selected at env creation.
+                img = self.env.render()
         img = img.astype(np.uint8)
         return img
 
+    def render(self, camera_name=None):
+        return self._get_image(camera_name=camera_name)
+
+    def _get_stateful_env(self):
+        for candidate in (self.env, getattr(self.env, "unwrapped", None)):
+            if candidate is not None:
+                has_get = hasattr(candidate, "get_env_state")
+                has_set = hasattr(candidate, "set_env_state")
+                if has_get and has_set:
+                    return candidate
+        return None
+
+    def get_env_state(self):
+        stateful_env = self._get_stateful_env()
+        if stateful_env is None:
+            raise AttributeError("Environment does not expose get_env_state/set_env_state")
+        return stateful_env.get_env_state()
+
+    def set_env_state(self, state):
+        stateful_env = self._get_stateful_env()
+        if stateful_env is None:
+            raise AttributeError("Environment does not expose get_env_state/set_env_state")
+        stateful_env.set_env_state(state)
+
+    def sync_from(self, other):
+        self.set_env_state(other.get_env_state())
+
     def reset(self, seed=None):
         obs, info = self.env.reset(seed=seed)
         state = self._extract_state(obs)
@@ -68,4 +105,4 @@ def step(self, action):
         return image, state, reward, done, info
 
     def close(self):
-        self.env.close()
\ No newline at end of file
+        self.env.close()
diff --git a/models/action_head_utils.py b/models/action_head_utils.py
new file mode 100644
index 0000000..1874391
--- /dev/null
+++ b/models/action_head_utils.py
@@ -0,0 +1,37 @@
+"""Shared action-head building blocks."""
+
+import torch.nn as nn
+
+
+class ResidualMLPBlock(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+
+    def forward(self, x):
+        return x + self.net(x)
+
+
+class ResidualActionMLP(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, hidden_dim: int, num_blocks: int = 3):
+        super().__init__()
+        self.input_proj = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.SiLU(),
+        )
+        self.blocks = nn.Sequential(*[ResidualMLPBlock(hidden_dim) for _ in range(num_blocks)])
+        self.out = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+
+    def forward(self, x):
+        x = self.input_proj(x)
+        x = self.blocks(x)
+        return self.out(x)
diff --git a/models/diffusion_head.py b/models/diffusion_head.py
index 045f158..ff855a9 100644
--- a/models/diffusion_head.py
+++ b/models/diffusion_head.py
@@ -5,6 +5,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from .action_head_utils import ResidualActionMLP
 
 
 @dataclass
@@ -14,6 +15,7 @@ class DiffusionConfig:
     beta_end: float = 1e-2
     action_dim: int = 4
     cond_dim: int = 128 # conditional input dim
+    hidden_dim: int = 256
 
 
 def make_beta_schedule(cfg: DiffusionConfig):
@@ -62,19 +64,16 @@ class ActionDenoiseModel(nn.Module):
     t:     (B,) integer timestep
     cond:  (B, cond_dim) fused VLA token
     """
-    def __init__(self, cfg: DiffusionConfig, time_emb_dim=32, hidden_dim=128):
+    def __init__(self, cfg: DiffusionConfig, time_emb_dim=32):
         super().__init__()
         self.cfg = cfg
         self.time_emb = SinusoidalTimeEmbedding(time_emb_dim)
 
         in_dim = cfg.action_dim + time_emb_dim + cfg.cond_dim
-
-        self.net = nn.Sequential(
-            nn.Linear(in_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, cfg.action_dim),
+        self.net = ResidualActionMLP(
+            in_dim=in_dim,
+            out_dim=cfg.action_dim,
+            hidden_dim=cfg.hidden_dim,
         )
 
     def forward(self, x_t, t, cond):
diff --git a/models/flow_matching_head.py b/models/flow_matching_head.py
new file mode 100644
index 0000000..0fa46d4
--- /dev/null
+++ b/models/flow_matching_head.py
@@ -0,0 +1,96 @@
+"""Flow-matching policy head for action generation."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+from .action_head_utils import ResidualActionMLP
+
+
+@dataclass
+class FlowMatchingConfig:
+    action_dim: int
+    cond_dim: int
+    t_embed_dim: int = 32
+    sample_steps: int = 32
+    hidden_dim: int = 256
+
+class SinusoidalTime(nn.Module):
+    """This is kept same as the diffusion time embedding."""
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, t):
+        half = self.dim // 2
+        freqs = torch.exp(torch.linspace(0, torch.log(torch.tensor(1000.0)), half, device=t.device))
+        args = t.unsqueeze(-1).float() * freqs.unsqueeze(0)
+        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
+        return emb
+
+class FlowMatchingModel(nn.Module):
+    """
+    Predict velocity field for flow matching: v(x_t, t | cond)
+    I LOVE THIS BLOG from Federico Sarrocco
+    https://federicosarrocco.com/blog/flow-matching
+    (hence, the code follows the same structure)
+    """
+
+    def __init__(self, cfg: FlowMatchingConfig):
+        super().__init__()
+        self.time_emb = SinusoidalTime(cfg.t_embed_dim)
+        in_dim = cfg.action_dim + cfg.t_embed_dim + cfg.cond_dim
+        self.net = ResidualActionMLP(
+            in_dim=in_dim,
+            out_dim=cfg.action_dim,
+            hidden_dim=cfg.hidden_dim,
+        )
+
+    def forward(self, x_t, t, cond):
+        t_emb = self.time_emb(t)
+        x = torch.cat([x_t, t_emb, cond], dim=-1)
+        v_pred = self.net(x)
+        return v_pred
+
+class FlowMatchingPolicyHead(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.model = FlowMatchingModel(cfg)
+
+    def loss(self, actions, cond):
+        """
+        Conditional flow matching with a linear interpolation path.
+
+        We sample a source point x_0 ~ N(0, I), set x_1 to the demonstrated
+        action, and train the model to predict the constant velocity field
+        along x_t = (1 - t) * x_0 + t * x_1.
+        """
+        B = actions.size(0)
+        t = torch.rand(B, device=actions.device)
+
+        x_0 = torch.randn_like(actions)
+        t_expanded = t.unsqueeze(-1)
+        x_t = (1.0 - t_expanded) * x_0 + t_expanded * actions
+        target_v = actions - x_0
+
+        v_pred = self.model(x_t, t, cond)
+        return F.mse_loss(v_pred, target_v)
+
+    @torch.no_grad()
+    def sample(self, cond, n_samples=None):
+        B = cond.size(0) if n_samples is None else n_samples
+        if cond.size(0) != B:
+            cond = cond.expand(B, -1)
+
+        # Start at the source distribution and integrate forward to t = 1.
+        x_t = torch.randn(B, self.cfg.action_dim, device=cond.device)
+        dt = 1.0 / self.cfg.sample_steps
+
+        for step in range(self.cfg.sample_steps):
+            t = torch.full((B,), step * dt, device=cond.device)
+            v = self.model(x_t, t, cond)
+            x_t = x_t + v * dt
+
+        return x_t
diff --git a/models/vision/VISION.md b/models/vision/VISION.md
new file mode 100644
index 0000000..ec7cfb9
--- /dev/null
+++ b/models/vision/VISION.md
@@ -0,0 +1 @@
+TODO: write a comprehensive doc to reproduce training, and testing using CLI commands 
diff --git a/models/vision/__init__.py b/models/vision/__init__.py
new file mode 100644
index 0000000..246f12a
--- /dev/null
+++ b/models/vision/__init__.py
@@ -0,0 +1,2 @@
+from .tinycnn import TinyCNNEncoder
+from .hf_vit import HFCLIPViT, HFSiglipViT
\ No newline at end of file
diff --git a/models/vision/hf_vit.py b/models/vision/hf_vit.py
new file mode 100644
index 0000000..7a517eb
--- /dev/null
+++ b/models/vision/hf_vit.py
@@ -0,0 +1,96 @@
+# models/vision/hf_vit.py
+from __future__ import annotations
+from typing import Literal
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .registry import VisionEncoder, VisionEncoderCfg, register_vision_encoder
+
+
+class HFViTVisionEncoder(VisionEncoder):
+    """
+    HuggingFace Vision encoder wrapper (CLIP / SigLIP).
+
+    Input:  x (B,3,H,W) float in [0,1] (or uint8 in [0,255])
+    Output: (B,D) pooled CLS token, projected to cfg.d_model if needed
+    """
+    def __init__(
+        self,
+        cfg: VisionEncoderCfg,
+        hf_kind: Literal["clip", "siglip"],
+        default_pretrained: str,
+        mean: tuple[float, float, float],
+        std: tuple[float, float, float],
+    ):
+        super().__init__(cfg)
+
+        try:
+            if hf_kind == "clip":
+                from transformers import CLIPVisionModel
+                model_id = cfg.pretrained or default_pretrained
+                self.backbone = CLIPVisionModel.from_pretrained(model_id)
+            elif hf_kind == "siglip":
+                from transformers import SiglipVisionModel
+                model_id = cfg.pretrained or default_pretrained
+                self.backbone = SiglipVisionModel.from_pretrained(model_id)
+            else:
+                raise ValueError(f"Unsupported hf_kind={hf_kind}")
+        except ImportError as e:
+            raise ImportError("Install transformers: pip install transformers") from e
+
+        # project to cfg.d_model
+        hidden = int(self.backbone.config.hidden_size)
+        out_dim = int(cfg.d_model)
+        self.proj = nn.Identity() if out_dim == hidden else nn.Linear(hidden, out_dim)
+
+        # normalize
+        self.register_buffer("mean", torch.tensor(mean, dtype=torch.float32).view(1, 3, 1, 1), persistent=False)
+        self.register_buffer("std", torch.tensor(std, dtype=torch.float32).view(1, 3, 1, 1), persistent=False)
+
+        self.image_size = int(cfg.image_size or getattr(self.backbone.config, "image_size", 224))
+
+        # cfg has trainable, not freeze
+        if not cfg.trainable:
+            for p in self.backbone.parameters():
+                p.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dtype == torch.uint8: x = x.float() / 255.0
+        else: x = x.float()
+
+        if x.shape[-2:] != (self.image_size, self.image_size):
+            x = F.interpolate(x, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)
+
+        x = (x - self.mean) / self.std # normalize
+        out = self.backbone(pixel_values=x, return_dict=True).last_hidden_state  # (B, N, hidden)
+        cls = out[:, 0]
+        return self.proj(cls)
+
+
+@register_vision_encoder("hf_clip_vit")
+class HFCLIPViT(HFViTVisionEncoder):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__(
+            cfg=cfg,
+            hf_kind="clip",
+            default_pretrained="openai/clip-vit-base-patch32",
+            # source for the following mean/std -- default values from HuggingFace documentation
+            # https://hf.co/docs/transformers/main/model_doc/clip#transformers.CLIPImageProcessor.image_mean
+            mean=(0.48145466, 0.4578275, 0.40821073), # CLIP mean
+            std=(0.26862954, 0.26130258, 0.27577711), # CLIP std
+            # i know, they look weird, but trust me they work
+        )
+
+
+@register_vision_encoder("hf_siglip_vit")
+class HFSiglipViT(HFViTVisionEncoder):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__(
+            cfg=cfg,
+            hf_kind="siglip",
+            default_pretrained="google/siglip-base-patch16-224",
+            mean=(0.5, 0.5, 0.5), # SigLIP mean
+            std=(0.5, 0.5, 0.5), # SigLIP std
+        )
diff --git a/models/vision/registry.py b/models/vision/registry.py
new file mode 100644
index 0000000..0207f75
--- /dev/null
+++ b/models/vision/registry.py
@@ -0,0 +1,46 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Dict, Type, Optional, List
+
+import torch
+import torch.nn as nn
+
+@dataclass
+class VisionEncoderCfg:
+    name : str = "tinycnn"
+    d_model : int = 128
+    pretrained : Optional[str] = None
+    trainable : bool = False
+    image_size : Optional[int] = None
+
+
+class VisionEncoder(nn.Module):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__()
+        self.cfg = cfg
+
+    @property
+    def d_model(self) -> int:
+        return self.cfg.d_model
+
+_REGISTRY: Dict[str, Type[VisionEncoder]] = {}
+
+def register_vision_encoder(name: str):
+    name = name.lower()
+
+    def decorator(cls: Type[VisionEncoder]):
+        if name in _REGISTRY:
+            raise ValueError(f"Vision encoder '{name}' already registered by {_REGISTRY[name]}.")
+        _REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+def available_vision_encoders() -> List[str]:
+    return list(_REGISTRY.keys())
+
+def build_vision_encoder(cfg: VisionEncoderCfg) -> VisionEncoder:
+    name = cfg.name.lower()
+    if name not in _REGISTRY:
+        raise ValueError(f"Unknown vision encoder '{name}'. Available: {available_vision_encoders()}")
+    return _REGISTRY[name](cfg)
\ No newline at end of file
diff --git a/models/vision/tinycnn.py b/models/vision/tinycnn.py
new file mode 100644
index 0000000..8598f65
--- /dev/null
+++ b/models/vision/tinycnn.py
@@ -0,0 +1,25 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .registry import VisionEncoder, VisionEncoderCfg, register_vision_encoder
+
+@register_vision_encoder("tinycnn")
+class TinyCNNEncoder(VisionEncoder):
+    def __init__(self, cfg: VisionEncoderCfg):
+        super().__init__(cfg)
+        d_model = cfg.d_model
+
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=5, stride=2, padding=2)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
+        self.proj = nn.Linear(128, d_model)
+        self.ln = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = x.mean(dim=[2, 3])
+        x = self.proj(x)
+        x = self.ln(x)
+        return x
diff --git a/models/vla_diffusion_policy.py b/models/vla_diffusion_policy.py
index 5f6e5ab..309bb11 100644
--- a/models/vla_diffusion_policy.py
+++ b/models/vla_diffusion_policy.py
@@ -1,29 +1,78 @@
 """VLA Diffusion Policy Model."""
 
+import torch
 import torch.nn as nn
-from .encoders import ImageEncoderTinyCNN, TextEncoderTinyGRU, StateEncoderMLP
+from .encoders import TextEncoderTinyGRU, StateEncoderMLP
 from .fusion import FusionMLP
+
+# action heads
 from .diffusion_head import DiffusionConfig, DiffusionPolicyHead
+from .flow_matching_head import FlowMatchingConfig, FlowMatchingPolicyHead
+
+from .vision.registry import VisionEncoderCfg, build_vision_encoder
+import models.vision
 
 
 class VLADiffusionPolicy(nn.Module):
-    def __init__(self, vocab_size, state_dim, action_dim,
-                 d_model=128, diffusion_T=16):
+    def __init__(
+        self,
+        vocab_size,
+        state_dim,
+        action_dim,
+        d_model=128,
+        diffusion_T=16,
+        vision_cfg: VisionEncoderCfg | None = None,
+        use_flow_matching=False,
+        obs_history_len=1,
+        action_head_hidden_dim=256,
+    ):
         super().__init__()
-        self.img_encoder = ImageEncoderTinyCNN(d_model=d_model)
+
+        # notice how I removed `ImageEncoderTinyCNN` import and now we rely on registry.
+
+        if vision_cfg is None:
+            vision_cfg = VisionEncoderCfg(name="tinycnn", d_model=d_model)
+
+        # vision encoder outputs d_model
+        vision_cfg.d_model = d_model
+        self.vision_cfg = vision_cfg
+        self.img_encoder = build_vision_encoder(vision_cfg)
+        self.obs_history_len = obs_history_len
+        if self.obs_history_len > 1:
+            self.img_history_proj = nn.Linear(obs_history_len * d_model, d_model)
+            self.img_history_ln = nn.LayerNorm(d_model)
+
         self.txt_encoder = TextEncoderTinyGRU(vocab_size=vocab_size, d_word=64, d_model=d_model)
-        self.state_encoder = StateEncoderMLP(state_dim=state_dim, d_model=d_model)
+        self.state_encoder = StateEncoderMLP(state_dim=state_dim * obs_history_len, d_model=d_model)
         self.fusion = FusionMLP(d_model=d_model)
 
-        cfg = DiffusionConfig(
-            T=diffusion_T,
-            action_dim=action_dim,
-            cond_dim=d_model,
-        )
-        self.diffusion_head = DiffusionPolicyHead(cfg)
+        # decide which action-head to use
+        self.use_flow_matching = use_flow_matching
+        if self.use_flow_matching:
+            fm_cfg = FlowMatchingConfig(
+                action_dim=action_dim,
+                cond_dim=d_model,
+                hidden_dim=action_head_hidden_dim,
+            )
+            self.policy_head = FlowMatchingPolicyHead(fm_cfg)
+        else:
+            cfg = DiffusionConfig(
+                T=diffusion_T,
+                action_dim=action_dim,
+                cond_dim=d_model,
+                hidden_dim=action_head_hidden_dim,
+            )
+            self.policy_head = DiffusionPolicyHead(cfg)
 
     def encode_obs(self, image, text_tokens, state):
-        img_token = self.img_encoder(image)  # (B, d_model)
+        if image.dim() == 5:
+            bsz, history_len, channels, height, width = image.shape
+            image = image.view(bsz * history_len, channels, height, width)
+            img_token = self.img_encoder(image).view(bsz, history_len, -1)
+            img_token = self.img_history_proj(img_token.flatten(start_dim=1))
+            img_token = self.img_history_ln(img_token)
+        else:
+            img_token = self.img_encoder(image)  # (B, d_model)
         txt_token = self.txt_encoder(text_tokens)  # (B, d_model)
         state_token = self.state_encoder(state)  # (B, d_model)
         fused_context = self.fusion(img_token, txt_token, state_token)
@@ -34,7 +83,7 @@ def loss(self, image, text_tokens, state, actions):
         Compute the loss of the diffusion policy head given the image, text tokens, state, and actions.
         """
         cond = self.encode_obs(image, text_tokens, state)
-        return self.diffusion_head.loss(actions, cond)
+        return self.policy_head.loss(actions, cond)
 
     def act(self, image, text_tokens, state):
         """
@@ -44,5 +93,4 @@ def act(self, image, text_tokens, state):
         returns: (B, action_dim)
         """
         cond = self.encode_obs(image, text_tokens, state)
-        actions = self.diffusion_head.sample(cond)
-        return actions
+        return self.policy_head.sample(cond)
diff --git a/scripts/ACTION_NORMALIZATION.md b/scripts/ACTION_NORMALIZATION.md
new file mode 100644
index 0000000..ca0537b
--- /dev/null
+++ b/scripts/ACTION_NORMALIZATION.md
@@ -0,0 +1,175 @@
+# Action Normalization And Clipping
+
+This note explains the action normalization upgrade added to mini-VLA, why it matters, what changed in code, and what to run now.
+
+## Why This Change Was Needed
+
+Before this change, the model was trained directly on raw environment actions.
+
+That is usually a poor default because:
+
+- different action dimensions can have different scales
+- optimization becomes unnecessarily harder
+- generated actions can drift outside the simulator's legal action range
+
+The new behavior is:
+
+- train on normalized actions
+- save action normalization statistics in the checkpoint
+- unnormalize predicted actions at inference time
+- clip final actions to the environment action bounds before stepping the env
+
+## What Changed
+
+### Training
+
+In [scripts/train.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/train.py):
+
+- the dataset now computes per-dimension `action_mean`
+- the dataset now computes per-dimension `action_std`
+- actions are normalized before being passed to the model
+- the checkpoint now stores:
+  - `action_stats.mean`
+  - `action_stats.std`
+  - `action_stats.eps`
+
+Formula used during training:
+
+```text
+normalized_action = (raw_action - action_mean) / action_std
+```
+
+`action_std` is clamped with a small epsilon so division stays numerically stable.
+
+### Inference
+
+In [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py):
+
+- checkpoints now load `action_stats` if present
+- model outputs are converted back to raw environment actions
+- actions are clipped to the env action space before `env.step(...)`
+
+Formula used during inference:
+
+```text
+raw_action = predicted_normalized_action * action_std + action_mean
+clipped_action = clip(raw_action, action_low, action_high)
+```
+
+### Environment Wrapper
+
+In [envs/metaworld_env.py](/Users/keivalya/Desktop/Projects/mini-vla/envs/metaworld_env.py):
+
+- the wrapper now exposes:
+  - `action_low`
+  - `action_high`
+
+These come from the simulator action space and are used for clipping.
+
+## Behavior With Old Checkpoints
+
+Old checkpoints do not contain `action_stats`.
+
+Current fallback behavior:
+
+- `action_mean = 0`
+- `action_std = 1`
+
+That keeps old checkpoints loadable, but it does **not** give you the benefit of action normalization. To actually use this upgrade, retrain the model.
+
+## Example
+
+Assume one action dimension has:
+
+```text
+action_mean = 0.25
+action_std = 0.50
+```
+
+If the model predicts:
+
+```text
+predicted_normalized_action = -1.20
+```
+
+Then inference converts it back to env space as:
+
+```text
+raw_action = (-1.20 * 0.50) + 0.25 = -0.35
+```
+
+If the environment bounds for that dimension are:
+
+```text
+action_low = -1.0
+action_high = 1.0
+```
+
+Then the clipped action stays:
+
+```text
+-0.35
+```
+
+If instead the reconstructed action were `1.7`, it would be clipped to `1.0`.
+
+## Files Changed
+
+- [scripts/train.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/train.py)
+- [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py)
+- [envs/metaworld_env.py](/Users/keivalya/Desktop/Projects/mini-vla/envs/metaworld_env.py)
+
+## What You Need To Run
+
+Because the checkpoint format changed, retrain first.
+
+### Train a New Checkpoint
+
+```bash
+python3 -m scripts.train \
+  --dataset-path data/push_v3.npz \
+  --epochs 50 \
+  --batch-size 64 \
+  --lr 1e-4 \
+  --d-model 128 \
+  --diffusion-T 16 \
+  --use-flow-matching \
+  --save-path checkpoints/flow_matching_model_norm.pt \
+  --device cpu
+```
+
+### Test the New Checkpoint
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_norm.pt \
+  --env-name push-v3 \
+  --policy-camera-name topview \
+  --video-camera-name corner2 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu \
+  --save-video \
+  --video-dir videos_showcase
+```
+
+## What You Should Expect In Logs
+
+During training, you should now see action statistics printed once at startup:
+
+```text
+[train] action_mean= [...]
+[train] action_std= [...]
+```
+
+During testing, you should now see:
+
+```text
+[test] action_low= [...]
+[test] action_high= [...]
+[test] action_mean= [...]
+[test] action_std= [...]
+```
+
+That confirms the checkpoint and environment bounds are being used correctly.
diff --git a/scripts/EVALUATION_METRICS.md b/scripts/EVALUATION_METRICS.md
new file mode 100644
index 0000000..2c857ec
--- /dev/null
+++ b/scripts/EVALUATION_METRICS.md
@@ -0,0 +1,169 @@
+# Aggregate Evaluation Metrics And Multi-Seed Runs
+
+This note documents the minimal evaluation upgrade added to [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py).
+
+## What Changed
+
+The evaluation script now supports:
+
+- aggregate metrics across all evaluation episodes
+- multi-seed test runs in a single command
+- per-episode success tracking
+- seed-specific video filenames so outputs do not overwrite each other
+
+This was implemented with minimal scope in [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py). No new script was added.
+
+## New CLI Behavior
+
+### Existing single-seed behavior still works
+
+You can still run:
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_norm.pt \
+  --env-name push-v3 \
+  --seed 42 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu
+```
+
+### New multi-seed behavior
+
+You can now pass multiple seeds with `--seeds`.
+
+Example:
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_norm.pt \
+  --env-name push-v3 \
+  --seeds 42 43 44 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu
+```
+
+If `--seeds` is provided, it overrides `--seed`.
+
+## Metrics Now Reported
+
+For each episode, the script now prints:
+
+- episode reward
+- episode length in steps
+- episode success
+
+At the end of the run, it prints aggregate statistics across all episodes from all seeds:
+
+- aggregated reward: mean, std, min, max
+- aggregated steps: mean, std, min, max
+- overall success rate
+- total number of evaluated episodes
+- list of seeds used
+
+## Success Definition
+
+Episode success is tracked from the environment `info` dict:
+
+```text
+success = max(success, int(info.get("success", 0)))
+```
+
+That means an episode is counted as successful if the environment reports success at any step during the rollout.
+
+## Video Naming Change
+
+When `--save-video` is used, output videos now include the seed in the filename.
+
+Before:
+
+```text
+push-v3_ep001.mp4
+```
+
+Now:
+
+```text
+push-v3_seed42_ep001.mp4
+```
+
+This prevents collisions during multi-seed runs.
+
+## Example Commands
+
+### Single-seed evaluation with saved videos
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_norm.pt \
+  --env-name push-v3 \
+  --seed 42 \
+  --policy-camera-name topview \
+  --video-camera-name corner2 \
+  --episodes 3 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu \
+  --save-video \
+  --video-dir videos_eval
+```
+
+### Multi-seed evaluation without videos
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_norm.pt \
+  --env-name push-v3 \
+  --seeds 42 43 44 45 46 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu
+```
+
+### Multi-seed evaluation with videos
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_norm.pt \
+  --env-name push-v3 \
+  --seeds 42 43 44 \
+  --policy-camera-name topview \
+  --video-camera-name corner2 \
+  --episodes 2 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu \
+  --save-video \
+  --video-dir videos_multi_seed
+```
+
+## Example Output
+
+Per-episode logs:
+
+```text
+[test] Seed 42 Episode 1/5: reward=7.321, steps=84, success=1
+[test] Seed 42 Episode 2/5: reward=6.908, steps=91, success=1
+```
+
+Aggregate summary:
+
+```text
+[test] Aggregated rewards: mean=6.944, std=0.412, min=6.201, max=7.321
+[test] Aggregated steps: mean=88.400, std=5.238, min=80.000, max=97.000
+[test] Success rate: 86.7%
+[test] Total episodes: 15 across seeds=[42, 43, 44]
+```
+
+## Files Changed
+
+- [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py)
+
+## Why This Is Useful
+
+A single rollout can be misleading. Multi-seed evaluation gives you a more defensible estimate of how stable the learned policy actually is, and aggregate metrics make it easier to compare checkpoints without inspecting videos one by one.
diff --git a/scripts/OBSERVATION_HISTORY.md b/scripts/OBSERVATION_HISTORY.md
new file mode 100644
index 0000000..c3130fe
--- /dev/null
+++ b/scripts/OBSERVATION_HISTORY.md
@@ -0,0 +1,182 @@
+# Observation History
+
+This note documents the minimal observation-history upgrade added to mini-VLA.
+
+## What Was Implemented
+
+Instead of giving the policy only the current image and current state, the model can now consume a short history of recent observations.
+
+This implementation uses:
+
+- stacked image history
+- stacked state history
+- the same current action target
+
+The goal is to give the policy short-term temporal context without redesigning the policy head.
+
+## Why Observation History Was Chosen
+
+Between the two options:
+
+- observation history
+- action chunking
+
+observation history is the smaller change in this codebase. It keeps the action head unchanged and only modifies data preparation plus input encoding.
+
+## What Changed
+
+### 1. Dataset collection now saves episode boundaries
+
+In [scripts/collect_data.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/collect_data.py):
+
+- each sample now stores an `episode_id`
+
+This is required so history stacking does not accidentally cross episode boundaries.
+
+### 2. Training supports `--obs-history-len`
+
+In [scripts/train.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/train.py):
+
+- added `--obs-history-len`
+- the dataset now builds a history window for images and states
+- the beginning of each episode is padded by repeating the earliest available observation in that episode
+- the checkpoint now stores `obs_history_len`
+
+### 3. The model now encodes stacked observations
+
+In [models/vla_diffusion_policy.py](/Users/keivalya/Desktop/Projects/mini-vla/models/vla_diffusion_policy.py):
+
+- image history is encoded frame-by-frame with the same vision encoder
+- the resulting per-frame image embeddings are concatenated and projected back to `d_model`
+- state history is concatenated and passed through the state encoder
+
+Default behavior is unchanged when `obs_history_len=1`.
+
+### 4. Inference now matches training
+
+In [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py):
+
+- the script reads `obs_history_len` from the checkpoint
+- it maintains rolling image and state buffers during evaluation
+- the model receives the same history length at test time that it was trained with
+
+## Important Compatibility Note
+
+Observation history greater than `1` requires datasets that include `episode_ids`.
+
+Older datasets may not have that field.
+
+Current behavior:
+
+- `obs_history_len=1` works with old datasets
+- `obs_history_len>1` raises an error if `episode_ids` are missing
+
+If that happens, regenerate the dataset with the updated [scripts/collect_data.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/collect_data.py).
+
+## Example
+
+If you set:
+
+```text
+obs_history_len = 4
+```
+
+then, for each training or evaluation step, the policy sees:
+
+- the last 4 images
+- the last 4 states
+
+At the first step of an episode, there is no past history yet, so the first observation is repeated to fill the buffer.
+
+So the first few episode steps look like:
+
+Step 0:
+
+```text
+[obs0, obs0, obs0, obs0]
+```
+
+Step 1:
+
+```text
+[obs0, obs0, obs0, obs1]
+```
+
+Step 2:
+
+```text
+[obs0, obs0, obs1, obs2]
+```
+
+This keeps history well-defined from the first step onward.
+
+## Files Changed
+
+- [scripts/collect_data.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/collect_data.py)
+- [scripts/train.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/train.py)
+- [models/vla_diffusion_policy.py](/Users/keivalya/Desktop/Projects/mini-vla/models/vla_diffusion_policy.py)
+- [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py)
+
+## What To Run
+
+### 1. Regenerate the dataset
+
+This step is required if your current dataset does not contain `episode_ids`.
+
+```bash
+python3 -m scripts.collect_data \
+  --env-name push-v3 \
+  --camera-name topview \
+  --episodes 100 \
+  --max-steps 100 \
+  --instruction "push the object to the goal" \
+  --output-path data/push_v3_history.npz
+```
+
+### 2. Train with observation history
+
+Example with 4-step history:
+
+```bash
+python3 -m scripts.train \
+  --dataset-path data/push_v3_history.npz \
+  --epochs 50 \
+  --batch-size 64 \
+  --lr 1e-4 \
+  --d-model 128 \
+  --diffusion-T 16 \
+  --obs-history-len 4 \
+  --use-flow-matching \
+  --save-path checkpoints/flow_matching_history4.pt \
+  --device cpu
+```
+
+### 3. Test the checkpoint
+
+No extra history flag is needed at test time. The script reads it from the checkpoint.
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_history4.pt \
+  --env-name push-v3 \
+  --seeds 42 43 44 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --policy-camera-name topview \
+  --video-camera-name corner2 \
+  --device cpu \
+  --save-video \
+  --video-dir videos_history4
+```
+
+## Minimal Design Choices
+
+This implementation intentionally does not:
+
+- change the action head
+- add a recurrent model
+- add attention over time
+- change the environment interface
+
+It only adds short observation history with the least amount of code movement needed to keep train and test aligned.
diff --git a/scripts/RUNNING.md b/scripts/RUNNING.md
new file mode 100644
index 0000000..954a97b
--- /dev/null
+++ b/scripts/RUNNING.md
@@ -0,0 +1,214 @@
+# Running mini-VLA
+
+This guide covers the exact commands to:
+
+- collect a dataset
+- train a model
+- evaluate a model
+- save videos from a different camera angle than the one used by the policy
+
+Run all commands from the repository root:
+
+```bash
+cd /Users/keivalya/Desktop/Projects/mini-vla
+```
+
+Use module mode (`python3 -m ...`) so imports like `from models...` work correctly.
+
+## 1. Dataset
+
+An existing dataset is already present at [data/push_v3.npz](/Users/keivalya/Desktop/Projects/mini-vla/data/push_v3.npz).
+
+If you want to regenerate it, run:
+
+```bash
+mkdir -p data
+
+python3 -m scripts.collect_data \
+  --env-name push-v3 \
+  --camera-name corner \
+  --episodes 100 \
+  --max-steps 100 \
+  --instruction "push the object to the goal" \
+  --output-path data/push_v3.npz
+```
+
+Notes:
+
+- `--camera-name` controls the image view stored in the dataset.
+- If you trained a model on top-view images, keep using top-view at inference time.
+
+## 2. Train Flow Matching
+
+The flow-matching bug was fixed in the code, so older flow-matching checkpoints should be considered stale and retrained.
+
+Train a new flow-matching checkpoint with:
+
+```bash
+mkdir -p checkpoints
+
+python3 -m scripts.train \
+  --dataset-path data/push_v3.npz \
+  --epochs 50 \
+  --batch-size 64 \
+  --lr 1e-4 \
+  --d-model 128 \
+  --diffusion-T 16 \
+  --use-flow-matching \
+  --save-path checkpoints/flow_matching_model_fixed.pt \
+  --device cpu
+```
+
+If CUDA is available, replace `--device cpu` with `--device cuda`.
+
+## 3. Evaluate Normally
+
+This runs the policy and saves videos from the same camera used for inference.
+
+```bash
+mkdir -p videos_fm_eval
+
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_fixed.pt \
+  --env-name push-v3 \
+  --policy-camera-name topview \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu \
+  --save-video \
+  --video-dir videos_fm_eval
+```
+
+## 4. Evaluate With a Showcase Camera
+
+This is the recommended setup when:
+
+- the policy should keep using the camera it was trained on
+- the saved MP4 should show the robot from a better angle for demos or social posts
+
+Example:
+
+```bash
+mkdir -p videos_showcase
+
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_fixed.pt \
+  --env-name push-v3 \
+  --policy-camera-name topview \
+  --video-camera-name corner2 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu \
+  --save-video \
+  --video-dir videos_showcase
+```
+
+Meaning of the camera flags:
+
+- `--policy-camera-name`: camera used to render the image that goes into the VLA
+- `--video-camera-name`: camera used only for the saved video
+- `--video-rotate`: optional override to rotate saved video frames by `0`, `90`, `180`, or `270` degrees
+
+If `--video-camera-name` is omitted, the saved video uses the same camera as the policy.
+By default, showcase videos recorded from a different camera than the policy camera are automatically rotated by `180` degrees.
+
+## 5. Camera Recommendations
+
+For model behavior:
+
+- use the same camera view the model was trained on
+
+For showcase videos:
+
+- `corner`
+- `corner2`
+- `corner3`
+- `corner4`
+- `behindGripper`
+
+Example combinations:
+
+- trained on `topview` -> infer with `--policy-camera-name topview`
+- save a nicer demo video -> add `--video-camera-name corner2`
+
+## 6. Common Errors
+
+### `ModuleNotFoundError: No module named 'models'`
+
+Cause:
+
+- running `python3 scripts/train.py` or `python3 scripts/test.py` directly
+
+Fix:
+
+```bash
+python3 -m scripts.train ...
+python3 -m scripts.test ...
+```
+
+Alternative:
+
+```bash
+PYTHONPATH=. python3 scripts/train.py ...
+PYTHONPATH=. python3 scripts/test.py ...
+```
+
+### Video still looks like the policy camera
+
+Cause:
+
+- some Meta-World/Gym render stacks only support the camera chosen when the environment is created
+
+Current behavior:
+
+- the code falls back to the policy camera instead of crashing
+
+## 7. Quick Copy-Paste Commands
+
+Train:
+
+```bash
+python3 -m scripts.train \
+  --dataset-path data/push_v3.npz \
+  --epochs 50 \
+  --batch-size 64 \
+  --lr 1e-4 \
+  --d-model 128 \
+  --diffusion-T 16 \
+  --use-flow-matching \
+  --save-path checkpoints/flow_matching_model_fixed.pt \
+  --device cpu
+```
+
+Test with same camera:
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_fixed.pt \
+  --env-name push-v3 \
+  --policy-camera-name topview \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu \
+  --save-video \
+  --video-dir videos_fm_eval
+```
+
+Test with separate showcase camera:
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_model_fixed.pt \
+  --env-name push-v3 \
+  --policy-camera-name topview \
+  --video-camera-name corner2 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --device cpu \
+  --save-video \
+  --video-dir videos_showcase
+```
diff --git a/scripts/STRONGER_ACTION_HEAD.md b/scripts/STRONGER_ACTION_HEAD.md
new file mode 100644
index 0000000..3d454e6
--- /dev/null
+++ b/scripts/STRONGER_ACTION_HEAD.md
@@ -0,0 +1,163 @@
+# Stronger Action Head
+
+This note documents the stronger action-head upgrade added to mini-VLA.
+
+## What Changed
+
+Previously, both policy heads used a very small MLP:
+
+- linear
+- ReLU
+- linear
+- ReLU
+- linear
+
+That is cheap, but it is also a bottleneck. The action head is where the model turns fused context into actual control outputs, so making it slightly stronger is a high-value change.
+
+The new action head is a shared residual MLP used by both:
+
+- diffusion action denoising
+- flow-matching velocity prediction
+
+## New Architecture
+
+The new shared head uses:
+
+- input projection
+- SiLU activation
+- multiple residual MLP blocks
+- LayerNorm before the final projection
+
+Concretely:
+
+```text
+input -> Linear -> SiLU -> Residual Blocks -> LayerNorm -> SiLU -> Linear -> action output
+```
+
+Each residual block is:
+
+```text
+x -> LayerNorm -> Linear -> SiLU -> Linear -> +x
+```
+
+This gives the head:
+
+- more capacity
+- better gradient flow than the old shallow MLP
+- stronger conditioning for the same encoders
+
+## Files Changed
+
+- [models/action_head_utils.py](/Users/keivalya/Desktop/Projects/mini-vla/models/action_head_utils.py)
+- [models/diffusion_head.py](/Users/keivalya/Desktop/Projects/mini-vla/models/diffusion_head.py)
+- [models/flow_matching_head.py](/Users/keivalya/Desktop/Projects/mini-vla/models/flow_matching_head.py)
+- [models/vla_diffusion_policy.py](/Users/keivalya/Desktop/Projects/mini-vla/models/vla_diffusion_policy.py)
+- [scripts/train.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/train.py)
+- [scripts/test.py](/Users/keivalya/Desktop/Projects/mini-vla/scripts/test.py)
+
+## New CLI Option
+
+Training now supports:
+
+```text
+--action-head-hidden-dim
+```
+
+This controls the hidden width of the stronger residual action head.
+
+Default:
+
+```text
+256
+```
+
+That value is saved in the checkpoint and automatically reused at test time.
+
+## Important Compatibility Note
+
+This upgrade changes the actual parameterization of the action head.
+
+That means older checkpoints trained with the old shallow head should be treated as incompatible for this feature. Retrain the model with the new code.
+
+## How To Train
+
+Example with flow matching, action normalization, and 4-step observation history:
+
+```bash
+python3 -m scripts.train \
+  --dataset-path data/push_v3_history.npz \
+  --epochs 50 \
+  --batch-size 64 \
+  --lr 1e-4 \
+  --d-model 128 \
+  --diffusion-T 16 \
+  --obs-history-len 4 \
+  --action-head-hidden-dim 256 \
+  --use-flow-matching \
+  --save-path checkpoints/flow_matching_history4_stronghead.pt \
+  --device cpu
+```
+
+If you want a larger head, increase it:
+
+```bash
+python3 -m scripts.train \
+  --dataset-path data/push_v3_history.npz \
+  --epochs 50 \
+  --batch-size 64 \
+  --lr 1e-4 \
+  --d-model 128 \
+  --diffusion-T 16 \
+  --obs-history-len 4 \
+  --action-head-hidden-dim 384 \
+  --use-flow-matching \
+  --save-path checkpoints/flow_matching_history4_head384.pt \
+  --device cpu
+```
+
+## How To Test
+
+You do not need to pass `--action-head-hidden-dim` at test time. The script reads it from the checkpoint automatically.
+
+Example:
+
+```bash
+python3 -m scripts.test \
+  --checkpoint checkpoints/flow_matching_history4_stronghead.pt \
+  --env-name push-v3 \
+  --seeds 42 43 44 \
+  --episodes 5 \
+  --max-steps 150 \
+  --instruction "push the object to the goal" \
+  --policy-camera-name topview \
+  --video-camera-name corner2 \
+  --device cpu \
+  --save-video \
+  --video-dir videos_stronghead
+```
+
+The test logs now print:
+
+```text
+[test] action_head_hidden_dim=256
+```
+
+That confirms the checkpoint is loading the stronger head configuration correctly.
+
+## When To Use This
+
+Use the stronger head when:
+
+- the model underfits even with decent encoders
+- actions look noisy or weakly conditioned
+- you want more capacity in the control head before changing the encoders
+
+## Minimal Design Choice
+
+This change intentionally does not:
+
+- introduce attention in the action head
+- add a transformer decoder
+- change the diffusion or flow-matching objective
+
+It only strengthens the final action network with a small residual MLP, which is the lowest-risk capacity upgrade in this codebase.
diff --git a/scripts/collect_data.py b/scripts/collect_data.py
index 0686d55..84b3a56 100644
--- a/scripts/collect_data.py
+++ b/scripts/collect_data.py
@@ -51,6 +51,7 @@ def main():
     states = []
     actions = []
     texts = []
+    episode_ids = []
 
     # fixed instruction for this dataset
     instruction = args.instruction
@@ -72,6 +73,7 @@ def main():
             states.append(state.copy())
             actions.append(np.asarray(action, dtype=np.float32).copy())
             texts.append(instruction)
+            episode_ids.append(ep)
 
             # step env
             obs, reward, truncate, terminate, info = env.step(action)
@@ -89,6 +91,7 @@ def main():
     images = np.stack(images, axis=0)   # (N, H, W, 3)
     states = np.stack(states, axis=0)   # (N, state_dim)
     actions = np.stack(actions, axis=0) # (N, action_dim)
+    episode_ids = np.asarray(episode_ids, dtype=np.int64)
 
     # tokenize instructions
     tokenizer = SimpleTokenizer(vocab=None)
@@ -106,6 +109,7 @@ def main():
         actions=actions,
         text_ids=text_ids,
         vocab=tokenizer.vocab,
+        episode_ids=episode_ids,
     )
 
     print("Saved Meta-World push dataset to", args.output_path)
@@ -113,6 +117,7 @@ def main():
     print("  states:", states.shape)
     print("  actions:", actions.shape)
     print("  text_ids:", text_ids.shape)
+    print("  episode_ids:", episode_ids.shape)
 
 
 if __name__ == "__main__":
diff --git a/scripts/test.py b/scripts/test.py
index c9f0b07..9b3f5ab 100644
--- a/scripts/test.py
+++ b/scripts/test.py
@@ -2,6 +2,7 @@
 
 import os
 import argparse
+from collections import deque
 import numpy as np
 import torch
 import imageio.v2 as imageio
@@ -9,6 +10,41 @@
 from envs.metaworld_env import MetaWorldMT1Wrapper
 from models.vla_diffusion_policy import VLADiffusionPolicy
 from utils.tokenizer import SimpleTokenizer
+from models.vision.registry import VisionEncoderCfg
+
+
+def normalize_action_stats(action_stats: dict | None, action_dim: int) -> tuple[np.ndarray, np.ndarray]:
+    if action_stats is None:
+        return np.zeros(action_dim, dtype=np.float32), np.ones(action_dim, dtype=np.float32)
+
+    mean = np.asarray(action_stats["mean"], dtype=np.float32)
+    std = np.asarray(action_stats["std"], dtype=np.float32)
+    std = np.clip(std, float(action_stats.get("eps", 1e-6)), None)
+    return mean, std
+
+
+def summarize_metric(values: list[float]) -> str:
+    arr = np.asarray(values, dtype=np.float32)
+    return f"mean={arr.mean():.3f}, std={arr.std():.3f}, min={arr.min():.3f}, max={arr.max():.3f}"
+
+
+def rotate_frame(frame: np.ndarray, degrees: int) -> np.ndarray:
+    if degrees % 360 == 0:
+        return frame
+    if degrees % 90 != 0:
+        raise ValueError(f"Rotation must be a multiple of 90 degrees, got {degrees}")
+    k = (degrees // 90) % 4
+    return np.rot90(frame, k=k).copy()
+
+
+def resolve_video_rotation(policy_camera_name: str, video_camera_name: str, requested_rotation: int | None) -> int:
+    if requested_rotation is not None:
+        return requested_rotation
+    # Showcase camera captures in this stack are typically upside down relative
+    # to the policy view, while the policy camera should stay untouched.
+    if video_camera_name != policy_camera_name:
+        return 180
+    return 0
 
 
 def parse_args():
@@ -26,11 +62,30 @@ def parse_args():
         default="push-v3",
         help="Meta-World MT1 task name, e.g. push-v3, reach-v3, pick-place-v3",
     )
+    parser.add_argument(
+        "--policy-camera-name",
+        type=str,
+        default="topview",
+        help="Camera used for policy observations at inference time",
+    )
+    parser.add_argument(
+        "--video-camera-name",
+        type=str,
+        default=None,
+        help="Optional camera used only for saved videos; defaults to the policy camera",
+    )
     parser.add_argument(
         "--seed",
         type=int,
         default=42,
-        help="Random seed for the environment",
+        help="Random seed for single-seed evaluation",
+    )
+    parser.add_argument(
+        "--seeds",
+        type=int,
+        nargs="+",
+        default=None,
+        help="Optional list of seeds for multi-seed evaluation; overrides --seed",
     )
     parser.add_argument(
         "--episodes",
@@ -67,12 +122,18 @@ def parse_args():
         default="videos",
         help="Directory to save videos (if --save-video is set)",
     )
+    parser.add_argument(
+        "--video-rotate",
+        type=int,
+        default=None,
+        help="Optional override to rotate saved video frames by 0, 90, 180, or 270 degrees",
+    )
 
     return parser.parse_args()
 
 
 def load_model_and_tokenizer(checkpoint_path: str, device: torch.device):
-    ckpt = torch.load(checkpoint_path, map_location=device)
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
 
     vocab = ckpt["vocab"]
     state_dim = ckpt["state_dim"]
@@ -80,6 +141,15 @@ def load_model_and_tokenizer(checkpoint_path: str, device: torch.device):
     d_model = ckpt["d_model"]
     diffusion_T = ckpt["diffusion_T"]
 
+    # load vision encoder config
+    vision_cfg = None
+    if "vision_cfg" in ckpt:
+        vision_cfg = VisionEncoderCfg(**ckpt["vision_cfg"])
+
+    use_flow_matching = ckpt.get("use_flow_matching", False)
+    obs_history_len = ckpt.get("obs_history_len", 1)
+    action_head_hidden_dim = ckpt.get("action_head_hidden_dim", 128)
+
     vocab_size = max(vocab.values()) + 1
 
     model = VLADiffusionPolicy(
@@ -88,6 +158,10 @@ def load_model_and_tokenizer(checkpoint_path: str, device: torch.device):
         action_dim=action_dim,
         d_model=d_model,
         diffusion_T=diffusion_T,
+        vision_cfg=vision_cfg,
+        use_flow_matching=use_flow_matching,
+        obs_history_len=obs_history_len,
+        action_head_hidden_dim=action_head_hidden_dim,
     ).to(device)
 
     model.load_state_dict(ckpt["model_state_dict"])
@@ -95,7 +169,9 @@ def load_model_and_tokenizer(checkpoint_path: str, device: torch.device):
 
     tokenizer = SimpleTokenizer(vocab=vocab)
 
-    return model, tokenizer
+    action_mean, action_std = normalize_action_stats(ckpt.get("action_stats"), action_dim)
+
+    return model, tokenizer, action_mean, action_std, obs_history_len, action_head_hidden_dim
 
 
 def main():
@@ -105,65 +181,128 @@ def main():
 
     # load model + tokenizer
     print(f"[test] Loading checkpoint from {args.checkpoint}")
-    model, tokenizer = load_model_and_tokenizer(args.checkpoint, device)
+    model, tokenizer, action_mean, action_std, obs_history_len, action_head_hidden_dim = load_model_and_tokenizer(args.checkpoint, device)
 
     # encode instruction
     instr_tokens = tokenizer.encode(args.instruction)
     text_ids = torch.tensor(instr_tokens, dtype=torch.long).unsqueeze(0).to(device)  # (1, T_text)
 
-    # environment
-    env = MetaWorldMT1Wrapper(
-        env_name=args.env_name,
-        seed=args.seed,
-        render_mode="rgb_array",
-        camera_name="topview",
+    video_camera_name = args.video_camera_name or args.policy_camera_name
+    video_rotation = resolve_video_rotation(
+        policy_camera_name=args.policy_camera_name,
+        video_camera_name=video_camera_name,
+        requested_rotation=args.video_rotate,
     )
-
-    print(f"[test] Meta-World MT1 env: {args.env_name}")
-    print(f"[test] state_dim={env.state_dim}, action_dim={env.action_dim}, obs_shape={env.obs_shape}")
-
     if args.save_video:
         os.makedirs(args.video_dir, exist_ok=True)
 
-    # evaluation
-    for ep in range(args.episodes):
-        img, state, info = env.reset()
-        step = 0
-        ep_reward = 0.0
-
-        frames = [img.copy()]
-
-        done = False
-        while not done and step < args.max_steps:
-            img_t = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0  # (1, 3, H, W)
-            state_t = torch.from_numpy(state).float().unsqueeze(0) # (1, state_dim)
-
-            img_t = img_t.to(device)
-            state_t = state_t.to(device)
-
-            # inference
-            with torch.no_grad():
-                action_t = model.act(img_t, text_ids, state_t)  # (1, action_dim)
-            action_np = action_t.squeeze(0).cpu().numpy()
-
-            # step environment
-            img, state, reward, done, info = env.step(action_np)
-            ep_reward += reward
-            step += 1
-
-            frames.append(img.copy())
-
-        print(f"[test] Episode {ep+1}/{args.episodes}: reward={ep_reward:.3f}, steps={step}")
-
-        # save video
-        if args.save_video:
-            video_path = os.path.join(args.video_dir, f"{args.env_name}_ep{ep+1:03d}.mp4")
-            with imageio.get_writer(video_path, fps=20) as writer:
-                for f in frames:
-                    writer.append_data(f)
-            print(f"[test] Saved video to {video_path}")
-
-    env.close()
+    eval_seeds = args.seeds if args.seeds is not None else [args.seed]
+    episode_rewards = []
+    episode_steps = []
+    episode_successes = []
+
+    for seed in eval_seeds:
+        env = MetaWorldMT1Wrapper(
+            env_name=args.env_name,
+            seed=seed,
+            render_mode="rgb_array",
+            camera_name=args.policy_camera_name,
+        )
+        video_env = None
+        if args.save_video and video_camera_name != args.policy_camera_name:
+            video_env = MetaWorldMT1Wrapper(
+                env_name=args.env_name,
+                seed=seed,
+                render_mode="rgb_array",
+                camera_name=video_camera_name,
+            )
+
+        print(f"[test] Meta-World MT1 env: {args.env_name}")
+        print(f"[test] seed={seed}, state_dim={env.state_dim}, action_dim={env.action_dim}, obs_shape={env.obs_shape}")
+        print(f"[test] obs_history_len={obs_history_len}")
+        print(f"[test] action_head_hidden_dim={action_head_hidden_dim}")
+        print(f"[test] policy_camera={args.policy_camera_name}, video_camera={video_camera_name}")
+        print(f"[test] video_rotate={video_rotation}")
+        print(f"[test] action_low={env.action_low.tolist()}")
+        print(f"[test] action_high={env.action_high.tolist()}")
+        print(f"[test] action_mean={action_mean.tolist()}")
+        print(f"[test] action_std={action_std.tolist()}")
+
+        for ep in range(args.episodes):
+            img, state, info = env.reset(seed=seed + ep)
+            step = 0
+            ep_reward = 0.0
+            ep_success = 0
+            img_history = deque([img.copy() for _ in range(obs_history_len)], maxlen=obs_history_len)
+            state_history = deque([state.copy() for _ in range(obs_history_len)], maxlen=obs_history_len)
+
+            if video_env is not None:
+                try:
+                    video_env.reset(seed=seed + ep)
+                    video_env.sync_from(env)
+                    frames = [rotate_frame(video_env.render().copy(), video_rotation)]
+                except AttributeError:
+                    print("[test] Warning: env state sync is unavailable; falling back to policy camera for video.")
+                    video_env.close()
+                    video_env = None
+                    frames = [rotate_frame(img.copy(), video_rotation)]
+            else:
+                frames = [rotate_frame(img.copy(), video_rotation)]
+
+            done = False
+            while not done and step < args.max_steps:
+                if obs_history_len > 1:
+                    img_np = np.stack(list(img_history), axis=0)  # (H, H_img, W_img, 3)
+                    img_t = torch.from_numpy(img_np).permute(0, 3, 1, 2).float().unsqueeze(0) / 255.0
+                else:
+                    img_t = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0
+                state_t = torch.from_numpy(np.concatenate(list(state_history), axis=0)).float().unsqueeze(0)
+
+                img_t = img_t.to(device)
+                state_t = state_t.to(device)
+
+                with torch.no_grad():
+                    action_t = model.act(img_t, text_ids, state_t)
+                action_np = action_t.squeeze(0).cpu().numpy()
+                action_np = action_np * action_std + action_mean
+                action_np = np.clip(action_np, env.action_low, env.action_high)
+
+                img, state, reward, done, info = env.step(action_np)
+                ep_reward += reward
+                ep_success = max(ep_success, int(info.get("success", 0)))
+                step += 1
+                img_history.append(img.copy())
+                state_history.append(state.copy())
+
+                if video_env is not None:
+                    video_env.sync_from(env)
+                    frames.append(rotate_frame(video_env.render().copy(), video_rotation))
+                else:
+                    frames.append(rotate_frame(img.copy(), video_rotation))
+
+            episode_rewards.append(ep_reward)
+            episode_steps.append(step)
+            episode_successes.append(ep_success)
+            print(
+                f"[test] Seed {seed} Episode {ep+1}/{args.episodes}: "
+                f"reward={ep_reward:.3f}, steps={step}, success={ep_success}"
+            )
+
+            if args.save_video:
+                video_path = os.path.join(args.video_dir, f"{args.env_name}_seed{seed}_ep{ep+1:03d}.mp4")
+                with imageio.get_writer(video_path, fps=20) as writer:
+                    for f in frames:
+                        writer.append_data(f)
+                print(f"[test] Saved video to {video_path}")
+
+        env.close()
+        if video_env is not None:
+            video_env.close()
+
+    print(f"[test] Aggregated rewards: {summarize_metric(episode_rewards)}")
+    print(f"[test] Aggregated steps: {summarize_metric(episode_steps)}")
+    print(f"[test] Success rate: {100.0 * float(np.mean(episode_successes)):.1f}%")
+    print(f"[test] Total episodes: {len(episode_rewards)} across seeds={eval_seeds}")
     print("[test] Done.")
 
 
diff --git a/scripts/train.py b/scripts/train.py
index 5ca41a5..94c00d4 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -7,10 +7,14 @@
 from torch.utils.data import Dataset, DataLoader
 
 from models.vla_diffusion_policy import VLADiffusionPolicy
+from models.vision.registry import VisionEncoderCfg
+
+
+ACTION_STD_EPS = 1e-6
 
 
 class TrainingDataset(Dataset):
-    def __init__(self, path, resize_to=64):
+    def __init__(self, path, resize_to=64, obs_history_len=1):
         data = np.load(path, allow_pickle=True)
         self.images = data["images"]             # (N, H, W, 3)
         self.states = data["states"]             # (N, state_dim)
@@ -18,6 +22,20 @@ def __init__(self, path, resize_to=64):
         self.text_ids = data["text_ids"]         # (N, T_text)
         self.vocab = data["vocab"].item() if data["vocab"].shape == () else data["vocab"]
         self.resize_to = resize_to
+        self.obs_history_len = obs_history_len
+        self.action_mean = self.actions.mean(axis=0).astype(np.float32)
+        self.action_std = self.actions.std(axis=0).astype(np.float32)
+        self.action_std = np.clip(self.action_std, ACTION_STD_EPS, None)
+        self.episode_ids = data["episode_ids"] if "episode_ids" in data else None
+        if self.obs_history_len > 1 and self.episode_ids is None:
+            raise ValueError(
+                "Observation history requires `episode_ids` in the dataset. "
+                "Regenerate the dataset with the updated scripts.collect_data."
+            )
+        if self.episode_ids is not None:
+            self.episode_start_idx = {}
+            for idx, episode_id in enumerate(self.episode_ids.tolist()):
+                self.episode_start_idx.setdefault(int(episode_id), idx)
 
         try:
             import cv2
@@ -28,14 +46,33 @@ def __init__(self, path, resize_to=64):
     def __len__(self):
         return self.images.shape[0]
 
-    def __getitem__(self, idx):
-        img = self.images[idx]  # (H, W, 3), uint8
-        if self.cv2 is not None and (img.shape[0] != self.resize_to or img.shape[1] != self.resize_to):
-            img = self.cv2.resize(img, (self.resize_to, self.resize_to))
+    def _history_indices(self, idx):
+        if self.obs_history_len == 1:
+            return [idx]
 
-        img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0  # (3, H, W)
-        state = torch.from_numpy(self.states[idx]).float()
-        action = torch.from_numpy(self.actions[idx]).float()
+        episode_id = int(self.episode_ids[idx])
+        episode_start = self.episode_start_idx[episode_id]
+        indices = []
+        for offset in range(self.obs_history_len):
+            hist_idx = max(episode_start, idx - (self.obs_history_len - 1 - offset))
+            indices.append(hist_idx)
+        return indices
+
+    def __getitem__(self, idx):
+        history_indices = self._history_indices(idx)
+        images = []
+        states = []
+        for hist_idx in history_indices:
+            img = self.images[hist_idx]  # (H, W, 3), uint8
+            if self.cv2 is not None and (img.shape[0] != self.resize_to or img.shape[1] != self.resize_to):
+                img = self.cv2.resize(img, (self.resize_to, self.resize_to))
+            images.append(torch.from_numpy(img).permute(2, 0, 1).float() / 255.0)
+            states.append(torch.from_numpy(self.states[hist_idx]).float())
+
+        img = torch.stack(images, dim=0) if self.obs_history_len > 1 else images[0]
+        state = torch.cat(states, dim=0)
+        normalized_action = (self.actions[idx] - self.action_mean) / self.action_std
+        action = torch.from_numpy(normalized_action).float()
         text_ids = torch.from_numpy(self.text_ids[idx]).long()
         return img, state, action, text_ids
 
@@ -50,10 +87,24 @@ def parse_args():
     parser.add_argument("--lr", type=float, default=1e-4)
     parser.add_argument("--d-model", type=int, default=128)
     parser.add_argument("--diffusion-T", type=int, default=16)
+    parser.add_argument("--obs-history-len", type=int, default=1)
+    parser.add_argument("--action-head-hidden-dim", type=int, default=256)
     parser.add_argument("--save-path", type=str,
                         default="checkpoints/model.pt")
     parser.add_argument("--device", type=str, default="cuda",
                         help="'cuda' or 'cpu'")
+
+    # vision command-line arguments
+    parser.add_argument("--vision-name", type=str, default="tinycnn",
+                    help="tinycnn | hf_clip_vit | hf_siglip_vit")
+    parser.add_argument("--vision-pretrained", type=str, default=None,
+                    help="HF model id override (optional)")
+    parser.add_argument("--vision-trainable", action="store_true",
+                    help="If set, HF backbone is trainable")
+    parser.add_argument("--vision-image-size", type=int, default=None,
+                    help="Override vision encoder image size (e.g. 224)")
+    parser.add_argument("--use-flow-matching", action="store_true", default=False,
+                    help="If set, use flow-matching instead of diffusion")
     return parser.parse_args()
 
 
@@ -62,17 +113,32 @@ def main():
     os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
 
     device = torch.device(args.device if torch.cuda.is_available() else "cpu")
-    dataset = TrainingDataset(args.dataset_path, resize_to=args.resize_to)
+    dataset = TrainingDataset(args.dataset_path, resize_to=args.resize_to, obs_history_len=args.obs_history_len)
     vocab_size = max(dataset.vocab.values()) + 1
     state_dim = dataset.states.shape[1]
     action_dim = dataset.actions.shape[1]
+    print("[train] action_mean=", dataset.action_mean.tolist())
+    print("[train] action_std=", dataset.action_std.tolist())
+
+    # vision config takes care of loading the correct vision encoder from CLI
+    vision_cfg = VisionEncoderCfg(
+        name=args.vision_name,
+        d_model=args.d_model,
+        pretrained=args.vision_pretrained,
+        trainable=args.vision_trainable,
+        image_size=args.vision_image_size,
+    )
 
     model = VLADiffusionPolicy(
         vocab_size=vocab_size,
         state_dim=state_dim,
         action_dim=action_dim,
         d_model=args.d_model,
-        diffusion_T=args.diffusion_T
+        diffusion_T=args.diffusion_T,
+        vision_cfg=vision_cfg,
+        use_flow_matching=args.use_flow_matching, # flow-matching / diffusion
+        obs_history_len=args.obs_history_len,
+        action_head_hidden_dim=args.action_head_hidden_dim,
     ).to(device)
 
     loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
@@ -107,6 +173,22 @@ def main():
             "action_dim": action_dim,
             "d_model": args.d_model,
             "diffusion_T": args.diffusion_T,
+            "use_flow_matching": args.use_flow_matching,
+            "obs_history_len": args.obs_history_len,
+            "action_head_hidden_dim": args.action_head_hidden_dim,
+            "action_stats": {
+                "mean": dataset.action_mean,
+                "std": dataset.action_std,
+                "eps": ACTION_STD_EPS,
+            },
+            # save vision encoder config
+            "vision_cfg": {
+                "name": vision_cfg.name,
+                "d_model": vision_cfg.d_model,
+                "pretrained": vision_cfg.pretrained,
+                "trainable": vision_cfg.trainable,
+                "image_size": vision_cfg.image_size,
+            },
         },
         args.save_path,
     )