Custom Models

Register a new preset

from vlarobot.models.base import VLAModel
from vlarobot.models.action import Action
from vlarobot.models.registry import MODEL_REGISTRY

class MyCustomVLA(VLAModel):
    def load(self):
        # Load your model weights
        pass

    def predict(self, image, instruction, **kwargs):
        # Run inference
        return Action(x=0.0, y=0.0, z=0.0, roll=0.0, pitch=0.0, yaw=0.0, gripper=1.0)

MODEL_REGISTRY["my-custom-vla"] = ("my_module:MyCustomVLA", "org/my-model")

Then use it like any other preset:

model = VLAModel.from_preset("my-custom-vla")

Load from checkpoint

model = VLAModel.from_checkpoint("./checkpoints/my-finetuned-model")