mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-25 18:56:09 +00:00
Khalil Meftah
8ff4a96b5e
- add isolated configs, models, processors, and tests - add shared distributional value utilities - add VF overfit comparison harness - add standalone Gemma3 VLM alignment workflow - keep the committed RECAP baseline unchanged
85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
import sys
|
|
from types import ModuleType, SimpleNamespace
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from lerobot.configs import FeatureType, PolicyFeature
|
|
from lerobot.rewards.factory import get_reward_model_class, make_reward_model_config
|
|
from lerobot.rewards.nanovlm_value_function.configuration_nanovlm_value_function import (
|
|
NanoVLMVFConfig,
|
|
)
|
|
from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
|
|
|
|
CAMERA = "observation.images.top"
|
|
|
|
|
|
def test_config_and_factory_registration():
|
|
config = make_reward_model_config("nanovlm_value_function")
|
|
assert isinstance(config, NanoVLMVFConfig)
|
|
assert get_reward_model_class("nanovlm_value_function").__name__ == "NanoVLMVFRewardModel"
|
|
|
|
|
|
def test_nanovlm_model_forward(monkeypatch):
|
|
from lerobot.rewards.nanovlm_value_function.modeling_nanovlm_value_function import (
|
|
NanoVLMVFRewardModel,
|
|
)
|
|
|
|
class FakeVision(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(1))
|
|
|
|
def forward(self, image):
|
|
return torch.ones(image.shape[0], 4, 6)
|
|
|
|
class FakeProjector(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.proj = nn.Linear(6, 8)
|
|
|
|
def forward(self, features):
|
|
return self.proj(features)
|
|
|
|
class FakeDecoder(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.token_embedding = nn.Embedding(100, 8)
|
|
|
|
def forward(self, inputs, attention_mask=None):
|
|
return inputs, None
|
|
|
|
class FakeNano(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.cfg = SimpleNamespace(lm_hidden_dim=8)
|
|
self.vision_encoder = FakeVision()
|
|
self.MP = FakeProjector()
|
|
self.decoder = FakeDecoder()
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, path):
|
|
return cls()
|
|
|
|
fake_module = ModuleType("models.vision_language_model")
|
|
fake_module.VisionLanguageModel = FakeNano
|
|
monkeypatch.setitem(sys.modules, "models.vision_language_model", fake_module)
|
|
|
|
config = NanoVLMVFConfig(
|
|
device="cpu",
|
|
nanovlm_code_path="third_party/nanoVLM",
|
|
)
|
|
config.input_features = {CAMERA: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16))}
|
|
model = NanoVLMVFRewardModel(config)
|
|
batch = {
|
|
CAMERA: torch.rand(1, 3, 16, 16),
|
|
CAMERA + ".mask": torch.ones(1, dtype=torch.bool),
|
|
OBS_LANGUAGE_TOKENS: torch.ones(1, 4, dtype=torch.long),
|
|
OBS_LANGUAGE_ATTENTION_MASK: torch.ones(1, 4, dtype=torch.bool),
|
|
"mc_return": torch.tensor([-0.5]),
|
|
"is_terminal": torch.tensor([False]),
|
|
}
|
|
loss, metrics = model(batch)
|
|
assert torch.isfinite(loss)
|
|
assert -1.0 <= metrics["predicted_value_mean"] <= 0.0
|