132 lines
5.4 KiB
Python
132 lines
5.4 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
DEFAULT_FRAME_SKIP = 4
|
|
|
|
MODEL_INPUT_SPEC: dict[str, tuple[tuple[int, ...], str]] = {
|
|
"img": ((1, 12, 128, 256), "uint8"),
|
|
"big_img": ((1, 12, 128, 256), "uint8"),
|
|
"desire_pulse": ((1, 25, 8), "float32"),
|
|
"traffic_convention": ((1, 2), "float32"),
|
|
"features_buffer": ((1, 24, 512), "float32"),
|
|
"action_t": ((1, 2), "float32"),
|
|
}
|
|
|
|
|
|
def spec_from_meta(meta: dict) -> dict[str, tuple[tuple[int, ...], str]] | None:
|
|
shapes = meta.get("input_shapes")
|
|
if not shapes:
|
|
return None
|
|
return {name: (tuple(shape), "uint8" if name in ("img", "big_img") else "float32")
|
|
for name, shape in shapes.items()}
|
|
|
|
|
|
class TemporalInputState:
|
|
def __init__(self, frame_skip: int, spec: dict[str, tuple[tuple[int, ...], str]] = MODEL_INPUT_SPEC):
|
|
self.frame_skip = frame_skip
|
|
img = spec["img"][0]
|
|
fb = spec["features_buffer"][0]
|
|
dp = spec["desire_pulse"][0]
|
|
|
|
self.n_frames = img[1] // 6
|
|
img_q_shape = (frame_skip * (self.n_frames - 1) + 1, 6, img[2], img[3])
|
|
self._img_shape = img
|
|
self._fb_shape = fb
|
|
self._dp_shape = dp
|
|
feat_dim = math.prod(fb[2:])
|
|
|
|
self.img_q = np.zeros(img_q_shape, dtype=np.uint8)
|
|
self.big_img_q = np.zeros(img_q_shape, dtype=np.uint8)
|
|
self.feat_q = np.zeros((frame_skip * fb[1], fb[0], feat_dim), dtype=np.float32)
|
|
self.desire_q = np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32)
|
|
self.prev_desire = np.zeros(dp[2], dtype=np.float32)
|
|
self.prev_feat = np.zeros((fb[0], feat_dim), dtype=np.float32)
|
|
|
|
@staticmethod
|
|
def _shift_append(q: np.ndarray, new_val: np.ndarray) -> None:
|
|
q[:-1] = q[1:]
|
|
q[-1] = new_val
|
|
|
|
def push_and_materialize(self, warped: np.ndarray, desire_pulse: np.ndarray,
|
|
traffic_convention: np.ndarray, action_t: np.ndarray,
|
|
) -> dict[str, np.ndarray]:
|
|
fs = self.frame_skip
|
|
|
|
cur = desire_pulse.astype(np.float32).copy()
|
|
cur[0] = 0
|
|
pulse = np.where(cur - self.prev_desire > 0.99, cur, 0).astype(np.float32)
|
|
self.prev_desire[:] = cur
|
|
|
|
self._shift_append(self.img_q, warped[0])
|
|
self._shift_append(self.big_img_q, warped[1])
|
|
self._shift_append(self.desire_q, pulse.reshape(self._dp_shape[0], self._dp_shape[2]))
|
|
self._shift_append(self.feat_q, self.prev_feat)
|
|
|
|
dp = self._dp_shape
|
|
return {
|
|
"img": np.ascontiguousarray(self.img_q[::fs]).reshape(self._img_shape),
|
|
"big_img": np.ascontiguousarray(self.big_img_q[::fs]).reshape(self._img_shape),
|
|
"features_buffer": np.ascontiguousarray(self.feat_q[::fs]).reshape(self._fb_shape),
|
|
"desire_pulse": self.desire_q.reshape(dp[1], fs, dp[0], dp[2]).max(axis=1).reshape(dp),
|
|
"traffic_convention": traffic_convention.astype(np.float32).reshape(1, -1),
|
|
"action_t": action_t.astype(np.float32).reshape(1, -1),
|
|
}
|
|
|
|
def note_hidden_state(self, model_output: np.ndarray, hidden_slice: slice) -> None:
|
|
self.prev_feat[:] = model_output[hidden_slice].reshape(self.prev_feat.shape)
|
|
|
|
|
|
class SplitTemporalState:
|
|
|
|
def __init__(self, frame_skip: int, img_shape: tuple[int, ...],
|
|
feature_shape: tuple[int, ...], desire_shape: tuple[int, ...]):
|
|
self.frame_skip = frame_skip
|
|
self._img_shape = tuple(img_shape)
|
|
self._fb_shape = tuple(feature_shape)
|
|
self._dp_shape = tuple(desire_shape)
|
|
|
|
n_frames = img_shape[1] // 6
|
|
img_q_shape = (frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3])
|
|
self.img_q = np.zeros(img_q_shape, dtype=np.uint8)
|
|
self.big_img_q = np.zeros(img_q_shape, dtype=np.uint8)
|
|
self.feat_q = np.zeros((frame_skip * (feature_shape[1] - 1) + 1, feature_shape[0], feature_shape[2]),
|
|
dtype=np.float32)
|
|
self.desire_q = np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), dtype=np.float32)
|
|
self.prev_desire = np.zeros(desire_shape[2], dtype=np.float32)
|
|
|
|
def materialize_vision(self, warped: np.ndarray, desire: np.ndarray) -> dict[str, np.ndarray]:
|
|
fs = self.frame_skip
|
|
cur = desire.astype(np.float32).copy()
|
|
cur[0] = 0
|
|
pulse = np.where(cur - self.prev_desire > 0.99, cur, 0).astype(np.float32)
|
|
self.prev_desire[:] = cur
|
|
|
|
TemporalInputState._shift_append(self.img_q, warped[0])
|
|
TemporalInputState._shift_append(self.big_img_q, warped[1])
|
|
TemporalInputState._shift_append(self.desire_q, pulse.reshape(self._dp_shape[0], self._dp_shape[2]))
|
|
return {
|
|
"img": np.ascontiguousarray(self.img_q[::fs]).reshape(self._img_shape),
|
|
"big_img": np.ascontiguousarray(self.big_img_q[::fs]).reshape(self._img_shape),
|
|
}
|
|
|
|
def materialize_policy(self, vision_feature: np.ndarray, traffic_convention: np.ndarray,
|
|
action_t: np.ndarray | None = None) -> dict[str, np.ndarray]:
|
|
fs = self.frame_skip
|
|
TemporalInputState._shift_append(self.feat_q, vision_feature.reshape(self._fb_shape[0], self._fb_shape[2]))
|
|
dp = self._dp_shape
|
|
out = {
|
|
"features_buffer": np.ascontiguousarray(self.feat_q[::fs]).reshape(self._fb_shape),
|
|
"desire_pulse": self.desire_q.reshape(dp[1], fs, dp[0], dp[2]).max(axis=1).reshape(dp),
|
|
"traffic_convention": traffic_convention.astype(np.float32).reshape(1, -1),
|
|
}
|
|
if action_t is not None:
|
|
out["action_t"] = action_t.astype(np.float32).reshape(1, -1)
|
|
return out
|