Control Generator

framework/data_generator/controls/control_generator.py

@@ -0,0 +1,104 @@
+import os
+import sys
+path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.insert(0, path)
+
+import numpy as np
+from carball.controls.controls import ControlsCreator
+from carball.json_parser.game import Game
+import carball
+from data_generator.base_generator import BaseDataGenerator
+from rlbot.utils.logging_utils import get_logger
+
+
+class ReplayControlGen(BaseDataGenerator):
+    """
+    Generates Controler data from a set of local replays.
+    Use method get_data() for Iterable
+    """
+
+    def __init__(self, filepaths: list):
+        # filepaths:list of paths to replay files (no file extension), player: player index, frames: number of frames per output
+        """
+        Arguments:
+        filepaths -- a list of paths to replay files, exclude the .replay file extensions
+        """
+        #super().__init__()
+        self.files = filepaths
+        self.file = 0
+        self.game = gameFromFile(self.files[self.file])
+        self.maxFrame = len(self.game.frames.time)
+        self.currentFrame = 1
+        self.player = 0
+
+    def initialize(self, **kwargs):
+        pass
+
+    def has_next(self):
+        if self.currentFrame >= self.maxFrame and self.player >= len( self.game.players ) - 1 and self.file >= len( self.files ) - 1:
+            return False
+        return True
+
+    def _next(self) -> list:
+        """
+        :return: list, [ Throttle, Steer, pitch, yaw, roll, jump, boost, handbrake ]
+        """
+        if self.currentFrame >= self.maxFrame:
+            self.player += 1
+            self.currentFrame = 1
+        if self.player >= len(self.game.players) and self.file < len(self.files):
+            self.file += 1
+            self.game = gameFromFile(self.files[self.file])
+            self.maxFrame = len(self.game.frames.time)
+            self.player = 0
+            self.currentFrame = 1
+        player = self.game.players[self.player]
+        c = getControls(player, self.currentFrame)
+        self.currentFrame += 1
+        return c
+
+
+def gameFromFile(replaypath: str) -> Game:
+    """
+    Instantiats a game object from a replay and creates controls data.
+    Creates json file representaion of replay.
+    Argument:
+    replaypath -- filepath to a replay, exclude .replay extension.
+    :return: Game
+    """
+    _json = carball.decompile_replay(replaypath + '.replay',
+                                     replaypath + '.json',
+                                     overwrite=True)
+
+    game = Game()
+    game.initialize(loaded_json=_json)
+
+    ControlsCreator().get_controls(game)
+    return game
+
+
+def getControls(player, frame: int) -> list:
+    """
+    Gets the controls of a player at a frame
+    Arguments:
+    player -- player object, game.players[p]
+    frame -- frame to get the controls from
+    :return: list of format [ Throttle, Steer, pitch, yaw, roll, jump, boost, handbrake ]
+    """
+    c = player.controls
+    throttle = c.throttle.get(frame)
+    steer = c.steer.get(frame)
+    pitch = c.pitch.get(frame)
+    yaw = c.yaw.get(frame)
+    roll = c.roll.get(frame)
+    jump = c.jump.get(frame)
+    boost = c.boost.get(frame)
+    handbrake = c.handbrake.get(frame)
+    controls = [throttle, steer, pitch, yaw, roll, jump, boost, handbrake]
+    for c in range(len(controls)):
+        control = controls[c]
+        if control == False or control == None or np.isnan(control):
+            controls[c] = 0
+        elif control == True:
+            controls[c] = 1
+    return controls