Improved the training loop to reflect actual control rate (20 Hz vs 500 Hz in simulation previously). Used AMD GPU for faster training and this is what it looks like (able to balance for ~6 seconds in simulation)