Spaces:
Running on Zero
Running on Zero
File size: 7,478 Bytes
80f7bd9 | 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 149 150 151 152 153 154 155 156 157 | """
DOOM-FlyWire Retinotopy & Compound Eye Encoder
==============================================
Pure neural retinotopy encoder mapping ViZDoom screens directly into
biological Drosophila optic lobe photoreceptor ensembles:
- Ommatidia array downsampling (32x32 compound eye grid).
- Spatial contrast & Elementary Motion (T4/T5 optical flow).
- Direct sensory current injection into Left & Right Optic Lobes.
- Continuous target salience projection into Lobula Columnar receptive fields.
- ZERO heuristic if-else overrides.
"""
import numpy as np
import torch
import torch.nn.functional as F
from typing import Tuple, Dict, Any, Optional
class RetinotopyEncoder:
def __init__(
self,
num_neurons: int = 139248,
optic_left_indices: np.ndarray = None,
optic_right_indices: np.ndarray = None,
t4t5_indices: np.ndarray = None,
eye_res: int = 32,
device: str = "cuda"
):
self.num_neurons = num_neurons
self.device = torch.device(device if torch.cuda.is_available() and device == "cuda" else "cpu")
self.eye_res = eye_res
self.prev_frame_gray = None
# Receptive field mapping onto biological optic lobe neurons
n_sensors = min(2048, len(optic_left_indices) // 10)
rng = np.random.RandomState(42)
self.idx_left = torch.from_numpy(rng.choice(optic_left_indices, n_sensors, replace=False)).to(self.device)
self.idx_right = torch.from_numpy(rng.choice(optic_right_indices, n_sensors, replace=False)).to(self.device)
if t4t5_indices is not None and len(t4t5_indices) > 0:
n_motion = min(2048, len(t4t5_indices))
self.idx_motion = torch.from_numpy(rng.choice(t4t5_indices, n_motion, replace=False)).to(self.device)
else:
self.idx_motion = None
print(f"[RetinotopyEncoder] Initialized with {n_sensors} sensors per hemifield, {eye_res}x{eye_res} compound eye.")
def reset(self):
"""Resets temporal motion integration buffer."""
self.prev_frame_gray = None
def encode(
self,
frame_rgb: np.ndarray,
labels: list = None,
params: dict = None,
contrast_gain: float = 2.0
) -> Tuple[torch.Tensor, Dict[str, Any]]:
"""
Pure neural retinotopy encoder:
- Downsamples RGB into 32x32 ommatidia array.
- Computes local spatial contrast and temporal frame-difference (Elementary Motion).
- Injects sensory currents into Left Optic, Right Optic, and T4/T5 Columnar Neurons.
- Maps detected visual targets (enemies) to biological Lobula Columnar (LC) hemifields.
"""
# Convert to HWC format if CHW
if frame_rgb.shape[0] == 3 and frame_rgb.ndim == 3:
frame_rgb = np.transpose(frame_rgb, (1, 2, 0))
# Downsample grayscale to compound eye ommatidia grid (eye_res x eye_res)
gray = (0.299 * frame_rgb[:, :, 0] + 0.587 * frame_rgb[:, :, 1] + 0.114 * frame_rgb[:, :, 2]) / 255.0
gray_tensor = torch.from_numpy(gray.astype(np.float32)).unsqueeze(0).unsqueeze(0).to(self.device)
ommatidia = F.interpolate(gray_tensor, size=(self.eye_res, self.eye_res), mode="area").squeeze()
# Compute Elementary Motion (temporal frame difference / optical flow)
if self.prev_frame_gray is not None:
motion_field = torch.abs(ommatidia - self.prev_frame_gray)
motion_intensity = motion_field.flatten()
motion_total = float(motion_field.mean().item())
else:
motion_intensity = torch.zeros_like(ommatidia.flatten())
motion_total = 0.0
self.prev_frame_gray = ommatidia.clone()
# Compute spatial contrast deviation (local center-surround contrast)
contrast_field = torch.abs(ommatidia - ommatidia.mean())
# Split ommatidia into Left Hemifield and Right Hemifield
mid_res = self.eye_res // 2
left_eye_contrast = contrast_field[:, :mid_res].flatten()
right_eye_contrast = contrast_field[:, mid_res:].flatten()
# Receptive field mapping to Left and Right Optic Lobes
n_left = len(self.idx_left)
l_photoreceptors = left_eye_contrast.repeat((n_left // len(left_eye_contrast)) + 1)[:n_left]
n_right = len(self.idx_right)
r_photoreceptors = right_eye_contrast.repeat((n_right // len(right_eye_contrast)) + 1)[:n_right]
# Initialize whole-brain sensory injection tensor
I_sensory = torch.zeros(self.num_neurons, device=self.device, dtype=torch.float32)
# Inject sensory currents via pure tensor operations
I_sensory[self.idx_left] = l_photoreceptors * (contrast_gain * 3.0)
I_sensory[self.idx_right] = r_photoreceptors * (contrast_gain * 3.0)
# Map optical motion to T4/T5 columnar motion detection neurons
if self.idx_motion is not None:
n_mot = len(self.idx_motion)
m_receptors = motion_intensity.repeat((n_mot // len(motion_intensity)) + 1)[:n_mot]
I_sensory[self.idx_motion] = m_receptors * (contrast_gain * 2.5)
# Target Tracking: Detect enemies on screen and stimulate corresponding visual hemifield
has_target = False
target_offset = 0.0
target_locked = False
enemy_count = 0
ENEMY_TYPES = {'MarineChainsawVzd', 'Demon', 'Cacodemon', 'HellKnight', 'DoomImp', 'BaronOfHell', 'LostSoul', 'ZombieMan', 'ShotgunGuy'}
if labels:
enemies = [l for l in labels if l.object_name in ENEMY_TYPES]
screen_w = frame_rgb.shape[1] if frame_rgb.ndim == 3 and frame_rgb.shape[0] != 3 else (frame_rgb.shape[2] if frame_rgb.ndim == 3 else 640)
center_x = screen_w / 2.0
enemy_count = len(enemies)
if enemies:
has_target = True
closest = min(enemies, key=lambda e: abs((e.x + e.width / 2.0) - center_x))
target_cx = closest.x + closest.width / 2.0
target_offset = (target_cx - center_x) / center_x
target_locked = abs(target_cx - center_x) < (screen_w * 0.08)
if target_locked:
# Bilateral surge on frontal target lock (both eyes stimulated)
I_sensory[self.idx_left] += contrast_gain * 3.5
I_sensory[self.idx_right] += contrast_gain * 3.5
elif target_offset < -0.05:
# Enemy is in Left Hemifield: stimulate Left Optic Lobe
I_sensory[self.idx_left] += (1.0 + abs(target_offset) * 2.0) * contrast_gain * 3.0
elif target_offset > 0.05:
# Enemy is in Right Hemifield: stimulate Right Optic Lobe
I_sensory[self.idx_right] += (1.0 + abs(target_offset) * 2.0) * contrast_gain * 3.0
visual_meta = {
"has_target": has_target,
"target_offset": round(target_offset, 3),
"target_locked": target_locked,
"enemy_count": enemy_count,
"optic_left_input": float(I_sensory[self.idx_left].mean().item()),
"optic_right_input": float(I_sensory[self.idx_right].mean().item()),
"motion_total": motion_total
}
return I_sensory, visual_meta
|