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..ef6d468 100644 --- a/envs/metaworld_env.py +++ b/envs/metaworld_env.py @@ -17,6 +17,7 @@ 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] @@ -49,11 +50,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 +103,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/flow_matching_head.py b/models/flow_matching_head.py new file mode 100644 index 0000000..3c729a5 --- /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 + + +@dataclass +class FlowMatchingConfig: + action_dim: int + cond_dim: int + t_embed_dim: int = 32 + sample_steps: int = 32 + +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, hidden_dim=128): + super().__init__() + self.time_emb = SinusoidalTime(cfg.t_embed_dim) + in_dim = cfg.action_dim + cfg.t_embed_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), + ) + + 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..5dcdea4 100644 --- a/models/vla_diffusion_policy.py +++ b/models/vla_diffusion_policy.py @@ -1,26 +1,52 @@ """VLA Diffusion Policy Model.""" 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, + ): 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.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.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) + self.policy_head = FlowMatchingPolicyHead(fm_cfg) + else: + cfg = DiffusionConfig(T=diffusion_T, action_dim=action_dim, cond_dim=d_model) + self.policy_head = DiffusionPolicyHead(cfg) def encode_obs(self, image, text_tokens, state): img_token = self.img_encoder(image) # (B, d_model) @@ -34,7 +60,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 +70,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/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/test.py b/scripts/test.py index c9f0b07..88666d7 100644 --- a/scripts/test.py +++ b/scripts/test.py @@ -9,6 +9,26 @@ 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 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,6 +46,18 @@ 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, @@ -67,6 +99,12 @@ 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() @@ -80,6 +118,13 @@ 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) + vocab_size = max(vocab.values()) + 1 model = VLADiffusionPolicy( @@ -88,6 +133,8 @@ 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, ).to(device) model.load_state_dict(ckpt["model_state_dict"]) @@ -116,11 +163,27 @@ def main(): env_name=args.env_name, seed=args.seed, render_mode="rgb_array", - camera_name="topview", + camera_name=args.policy_camera_name, + ) + 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, ) + video_env = None + if args.save_video and video_camera_name != args.policy_camera_name: + video_env = MetaWorldMT1Wrapper( + env_name=args.env_name, + seed=args.seed, + render_mode="rgb_array", + camera_name=video_camera_name, + ) 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}") + print(f"[test] policy_camera={args.policy_camera_name}, video_camera={video_camera_name}") + print(f"[test] video_rotate={video_rotation}") if args.save_video: os.makedirs(args.video_dir, exist_ok=True) @@ -131,7 +194,17 @@ def main(): step = 0 ep_reward = 0.0 - frames = [img.copy()] + if video_env is not None: + try: + 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: @@ -151,7 +224,11 @@ def main(): ep_reward += reward step += 1 - frames.append(img.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)) print(f"[test] Episode {ep+1}/{args.episodes}: reward={ep_reward:.3f}, steps={step}") @@ -164,6 +241,8 @@ def main(): print(f"[test] Saved video to {video_path}") env.close() + if video_env is not None: + video_env.close() print("[test] Done.") diff --git a/scripts/train.py b/scripts/train.py index 5ca41a5..5a5062c 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -7,6 +7,7 @@ from torch.utils.data import Dataset, DataLoader from models.vla_diffusion_policy import VLADiffusionPolicy +from models.vision.registry import VisionEncoderCfg class TrainingDataset(Dataset): @@ -54,6 +55,18 @@ def parse_args(): 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() @@ -67,12 +80,23 @@ def main(): state_dim = dataset.states.shape[1] action_dim = dataset.actions.shape[1] + # 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 ).to(device) loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) @@ -107,6 +131,15 @@ def main(): "action_dim": action_dim, "d_model": args.d_model, "diffusion_T": args.diffusion_T, + "use_flow_matching": args.use_flow_matching, + # 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, )