| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import math
|
|
|
| ACTION_TYPES = [
|
| "CLICK",
|
| "DOUBLE_CLICK",
|
| "RIGHT_CLICK",
|
| "TYPE",
|
| "SCROLL_UP",
|
| "SCROLL_DOWN",
|
| "HOVER",
|
| "WAIT"
|
| ]
|
| NUM_ACTIONS = len(ACTION_TYPES)
|
|
|
| def tokenize(text: str, max_len: int = 32) -> torch.Tensor:
|
| tokens = [ord(c) % 256 for c in text.lower()[:max_len]]
|
| tokens += [0] * (max_len - len(tokens))
|
| return torch.tensor([tokens], dtype=torch.long)
|
|
|
|
|
| class InstructionEncoder(nn.Module):
|
| def __init__(self, vocab_size=2048, embed_dim=256):
|
| super().__init__()
|
| self.embedding = nn.Embedding(vocab_size, embed_dim)
|
| self.proj = nn.Sequential(
|
| nn.Linear(embed_dim, embed_dim),
|
| nn.GELU(),
|
| nn.Linear(embed_dim, embed_dim)
|
| )
|
|
|
| def forward(self, token_ids):
|
| embeds = self.embedding(token_ids)
|
| pooled = embeds.mean(dim=1)
|
| return self.proj(pooled)
|
|
|
|
|
| class AevanisOperatorHead(nn.Module):
|
| def __init__(self, visual_dim=256, hidden_dim=256, num_actions=NUM_ACTIONS):
|
| super().__init__()
|
| self.visual_dim = visual_dim
|
| self.hidden_dim = hidden_dim
|
|
|
|
|
| self.q_proj = nn.Linear(hidden_dim, hidden_dim)
|
| self.k_proj = nn.Linear(visual_dim, hidden_dim)
|
| self.v_proj = nn.Linear(visual_dim, hidden_dim)
|
| self.scale = 1.0 / math.sqrt(hidden_dim)
|
|
|
|
|
| self.heatmap_head = nn.Sequential(
|
| nn.Linear(hidden_dim, 64),
|
| nn.GELU(),
|
| nn.Linear(64, 1)
|
| )
|
|
|
|
|
| self.action_head = nn.Sequential(
|
| nn.Linear(hidden_dim * 2, hidden_dim),
|
| nn.LayerNorm(hidden_dim),
|
| nn.GELU(),
|
| nn.Linear(hidden_dim, num_actions)
|
| )
|
|
|
|
|
| self.coord_head = nn.Sequential(
|
| nn.Linear(hidden_dim * 2, hidden_dim),
|
| nn.LayerNorm(hidden_dim),
|
| nn.GELU(),
|
| nn.Linear(hidden_dim, 64),
|
| nn.GELU(),
|
| nn.Linear(64, 2),
|
| nn.Sigmoid()
|
| )
|
|
|
| def forward(self, visual_tokens, instr_embed):
|
| B, N, _ = visual_tokens.shape
|
| q = self.q_proj(instr_embed).unsqueeze(1)
|
| k = self.k_proj(visual_tokens)
|
| v = self.v_proj(visual_tokens)
|
|
|
| attn_scores = torch.bmm(q, k.transpose(1, 2)) * self.scale
|
| attn_weights = F.softmax(attn_scores, dim=-1)
|
| context = torch.bmm(attn_weights, v).squeeze(1)
|
|
|
| heatmap_logits = self.heatmap_head(k).squeeze(-1)
|
| heatmap_2d = heatmap_logits.reshape(B, 14, 14)
|
|
|
| multimodal = torch.cat([context, instr_embed], dim=-1)
|
| action_logits = self.action_head(multimodal)
|
| pred_coords = self.coord_head(multimodal)
|
|
|
| action_idx = torch.argmax(action_logits, dim=-1)
|
| action_names = [ACTION_TYPES[idx.item()] for idx in action_idx]
|
|
|
| return {
|
| "coords": pred_coords,
|
| "action_logits": action_logits,
|
| "action_name": action_names,
|
| "heatmap": heatmap_2d,
|
| "confidence": torch.softmax(action_logits, dim=-1).max(dim=-1)[0]
|
| }
|
|
|
|
|
| class AevanisOperatorAgent(nn.Module):
|
| """Standalone Agent with self-contained head and optional backbone loading."""
|
| def __init__(self, model_scale="nano", head_weights=None, backbone=None):
|
| super().__init__()
|
| self.model_scale = model_scale
|
| visual_dim = 256 if model_scale == "nano" else 384
|
| self.instr_encoder = InstructionEncoder(vocab_size=2048, embed_dim=256)
|
| self.operator_head = AevanisOperatorHead(visual_dim=visual_dim, hidden_dim=256)
|
| self.backbone = backbone
|
|
|
| if head_weights is not None:
|
| self.load_head(head_weights)
|
|
|
| def load_head(self, path: str):
|
| state = torch.load(path, map_location="cpu")
|
| cleaned = {k.replace("operator_head.", ""): v for k, v in state.items()}
|
| self.operator_head.load_state_dict(cleaned, strict=False)
|
| self.eval()
|
|
|
| def forward(self, screen_input, token_ids):
|
|
|
|
|
|
|
| if screen_input.dim() == 4:
|
| if self.backbone is not None:
|
| with torch.no_grad():
|
| features = self.backbone(screen_input)
|
| visual_tokens = features["spatial_tokens"]
|
| else:
|
|
|
| B = screen_input.shape[0]
|
| visual_tokens = torch.zeros(B, 196, self.operator_head.visual_dim)
|
| else:
|
| visual_tokens = screen_input
|
|
|
| instr_embed = self.instr_encoder(token_ids)
|
| return self.operator_head(visual_tokens, instr_embed)
|
|
|
| @classmethod
|
| def from_pretrained(cls, model_scale="nano", head_weights=None):
|
| return cls(model_scale=model_scale, head_weights=head_weights)
|
|
|
