Control Generator

framework/data_generator/controls/control_generator.py

@@ -0,0 +1,79 @@
+from rlbot.utils.logging_utils import get_logger
+from framework.data_generator.base_generator import BaseDataGenerator
+import carball
+from carball.json_parser.game import Game
+from carball.controls.controls import ControlsCreator
+import numpy as np
+
+class ReplayControlGen( BaseDataGenerator ):
+    def __init__(self, filepaths, player = 0, frames = 15 ):
+        # filepaths:list of paths to replay files (no file extension), player: player index, frames: number of frames per output
+        super().__init__( )
+        self.files = filepaths
+        self.file = 0
+        self.game = gameFromFile( self.files[self.file] )
+        self.maxFrame = len( self.game.frames.time )
+        self.frames = frames
+        self.currentFrame = 1
+        self.player = 0
+
+    def initialize( self, **kwargs ):
+        pass
+
+    def has_next(self):
+        if self.file < len( self.files ) - 1:
+            return True
+        if self.currentFrame + self.frames > self.maxFrame or self.player >= len( self.game.players ):
+            return False
+        return True
+
+    def _next(self) -> list:
+        if self.currentFrame + self.frames > self.maxFrame:
+            self.player += 1
+        if self.player >= len( self.game.players ):
+            self.file += 1
+            self.game = gameFromFile( self.files[self.file] )
+            self.maxFrame = len( self.game.frames.time )
+            self.player = 0
+        player = self.game.players[ self.player ]
+        startFrame = self.currentFrame
+        endFrame = startFrame + self.frames
+        self.currentFrame = endFrame
+        return sampleFrames( self.game, player, startFrame, endFrame )
+
+def gameFromFile( replayName ): #Do not include file extension
+    _json = carball.decompile_replay(replayName + '.replay',
+                                    replayName + '.json',
+                                    overwrite=True)
+
+    game = Game()
+    game.initialize(loaded_json=_json)
+
+    ControlsCreator().get_controls( game )
+    return game
+
+def sampleFrames( game, player, start, end ):
+    frames = []
+    for i in range( start, end ):
+        frames.append( getControls( player, i ) )
+    return frames
+
+def getControls( player, frame ):
+    #[ 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