Add Nexto & its toxic variant

Author: VirxEC

tests/atba/atba.bot.toml

@@ -21,7 +21,7 @@ fun_fact = "This is a test bot"
 source_link = "https://github.com/RLBot/RLBot"
 developer = "BotMaker"
 language = "Python 3"
-# ALL POSSIBLE TAGS: 1v1, teamplay, goalie, hoops, dropshot, snow-day, spike-rush, heatseaker, memebot
+# ALL POSSIBLE TAGS: 1v1, teamplay, goalie, hoops, dropshot, snow-day, spike-rush, heatseeker, memebot
 # NOTE: Only add the goalie tag if your bot only plays as a goalie; this directly contrasts with the teamplay tag!
 # NOTE: Only add a tag for a special game mode if you bot properly supports it
 tags = []

tests/nexto/agent.py

@@ -0,0 +1,88 @@
+import math
+import os
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.distributions import Categorical
+
+
+def make_lookup_table():
+    actions = []
+    # Ground
+    for throttle in (-1, 0, 1):
+        for steer in (-1, 0, 1):
+            for boost in (0, 1):
+                for handbrake in (0, 1):
+                    if boost == 1 and throttle != 1:
+                        continue
+                    actions.append(
+                        [throttle or boost, steer, 0, steer, 0, 0, boost, handbrake]
+                    )
+    # Aerial
+    for pitch in (-1, 0, 1):
+        for yaw in (-1, 0, 1):
+            for roll in (-1, 0, 1):
+                for jump in (0, 1):
+                    for boost in (0, 1):
+                        if jump == 1 and yaw != 0:  # Only need roll for sideflip
+                            continue
+                        if pitch == roll == jump == 0:  # Duplicate with ground
+                            continue
+                        # Enable handbrake for potential wavedashes
+                        handbrake = jump == 1 and (pitch != 0 or yaw != 0 or roll != 0)
+                        actions.append(
+                            [boost, yaw, pitch, yaw, roll, jump, boost, handbrake]
+                        )
+    actions = np.array(actions)
+    return actions
+
+
+class Agent:
+    _lookup_table = make_lookup_table()
+    state = None
+
+    def __init__(self):
+        cur_dir = os.path.dirname(os.path.realpath(__file__))
+        with open(os.path.join(cur_dir, "nexto-model.pt"), "rb") as f:
+            self.actor = torch.jit.load(f)
+        torch.set_num_threads(1)
+
+    def act(self, state, beta):
+        state = tuple(torch.from_numpy(s).float() for s in state)
+
+        with torch.no_grad():
+            out, weights = self.actor(state)
+        self.state = state
+
+        out = (out,)
+        max_shape = max(o.shape[-1] for o in out)
+        logits = torch.stack(
+            [
+                (
+                    l
+                    if l.shape[-1] == max_shape
+                    else F.pad(l, pad=(0, max_shape - l.shape[-1]), value=float("-inf"))
+                )
+                for l in out
+            ],
+            dim=1,
+        )
+
+        # beta = 0.5
+        if beta == 1:
+            actions = np.argmax(logits, axis=-1)
+        elif beta == -1:
+            actions = np.argmin(logits, axis=-1)
+        else:
+            if beta == 0:
+                logits[torch.isfinite(logits)] = 0
+            else:
+                logits *= math.log((beta + 1) / (1 - beta), 3)
+            dist = Categorical(logits=logits)
+            actions = dist.sample()
+
+        # print(Categorical(logits=logits).sample())
+        parsed = self._lookup_table[actions.numpy().item()]
+
+        return parsed, weights