File size: 5,253 Bytes
d0044a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | 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
# Cross-Attention
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)
# 14x14 Heatmap
self.heatmap_head = nn.Sequential(
nn.Linear(hidden_dim, 64),
nn.GELU(),
nn.Linear(64, 1)
)
# Action Classifier
self.action_head = nn.Sequential(
nn.Linear(hidden_dim * 2, hidden_dim),
nn.LayerNorm(hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, num_actions)
)
# High-Precision (X, Y) Coordinates (sub-pixel sigmoid)
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):
# screen_input can be:
# 1. 14x14 Spatial Tokens [B, 196, D] directly
# 2. Raw RGB Image [B, 3, 224, 224] if backbone is attached
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:
# Standalone fallback: dummy projection for self-contained testing
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)
|