diff --git a/Assets/Scripts/ML Cozmo/CozmoAgent.cs b/Assets/Scripts/ML Cozmo/CozmoAgent.cs
index fe3339e..4e3f44a 100644
--- a/Assets/Scripts/ML Cozmo/CozmoAgent.cs	
+++ b/Assets/Scripts/ML Cozmo/CozmoAgent.cs	
@@ -31,8 +31,12 @@ namespace Cozmo
         public Camera renderCamera;
         [Tooltip("Reference to the CozmoMovement script")]
         public CozmoMovementController movementController;
+        [Tooltip("The time between decisions at inference")]
         public float timeBetweenDecisionsAtInference;
+        [Tooltip("The amout of actions that should be remembered by the agent.")]
         public int maxStoredMovementStates = 1;
+        [Tooltip("If activated the maxStoredMovementStates has no impact on the training.")]
+        public bool useOriginalActionMasking = true;
 
         private Academy academy;                        // CozmoAcademy
         private float timeSinceDecision;                // time since last decision
@@ -40,13 +44,13 @@ namespace Cozmo
         private int nearAreaLimit = 0;                  // X coordinate limit for the near to the imagecenter area 
         private int centerOfImageX = 0;                 // Middle of the image in x direction
         private MovementState lastChosenMovement = MovementState.Stop;  // The last action/movement that was executed 
+        private Queue<MovementState> lastActions = new Queue<MovementState>(); // Queue to store the last chosen Actions
 
-        private Queue<MovementState> lastActions = new Queue<MovementState>();
-
-        private double startTime = Time.time;
+        private double startTime = 0;
 
         private void Start()
         {
+            startTime = Time.time;
             academy = FindObjectOfType(typeof(CozmoAcademy)) as CozmoAcademy;
             imageProcessor = renderCamera.GetComponent<ImageProcessor>();
             nearAreaLimit = (int)(renderCamera.targetTexture.width / 2 * NEAR_AREA_PERCENTAGE_OFFSET);
@@ -66,50 +70,56 @@ namespace Cozmo
         }
 
         // Set ActionMask for training
+        // Needs to be called from CollectObservations()
         private void SetMask()
         {
-            // Do not allow stop decision after a stop 
-            if (lastChosenMovement == MovementState.Stop)
+            if (useOriginalActionMasking)
             {
-                SetActionMask(STOP);
+                // Stop is never allowed 
+                switch (lastChosenMovement)
+                {
+                    // Do not allow stop decision after a stop 
+                    case (MovementState.Stop):
+                        SetActionMask(STOP);
+                        break;
+                    // Do not allow stop after forward
+                    case (MovementState.Forward):
+                        SetActionMask(STOP);
+                        break;
+                    // Do not allow stop & left after right
+                    case (MovementState.Right):
+                        SetActionMask(STOP);
+                        SetActionMask(LEFT);
+                        break;
+                    // Do not allow stop & right after left
+                    case (MovementState.Left):
+                        SetActionMask(STOP);
+                        SetActionMask(RIGHT);
+                        break;
+                    default:
+                        throw new ArgumentException("Invalid MovementState.");
+                }
             }
-
-            // Do not allow left decision if right was in the last actions
-            if (lastActions.Contains(MovementState.Right))
+            else
             {
-                SetActionMask(LEFT);
-            }
+                // Do not allow stop decision after a stop 
+                if (lastChosenMovement == MovementState.Stop)
+                {
+                    SetActionMask(STOP);
+                }
 
-            // Do not allow right decision if left was in the last actions
-            if (lastActions.Contains(MovementState.Left))
-            {
-                SetActionMask(RIGHT);
-            }
+                // Do not allow left decision if right was in the last actions
+                if (lastActions.Contains(MovementState.Right))
+                {
+                    SetActionMask(LEFT);
+                }
 
-            //switch (lastChosenMovement)
-            //{
-            //    // Do not allow stop decision after a stop 
-            //    case (MovementState.Stop):
-            //        SetActionMask(STOP);
-            //        break;
-            //    // Do not allow stop after forward
-            //    case (MovementState.Forward):
-            //        //SetActionMask(STOP);
-            //        break;
-            //    // Do not allow stop & left after right
-            //    case (MovementState.Right):
-            //        //SetActionMask(STOP);
-            //        if (lastActions.Contains(MovementState.))
-            //            SetActionMask(LEFT);
-            //        break;
-            //    // Do not allow stop & right after left
-            //    case (MovementState.Left):
-            //        //SetActionMask(STOP);
-            //        SetActionMask(RIGHT);
-            //        break;
-            //    default:
-            //        throw new ArgumentException("Invalid MovementState.");
-            //}
+                // Do not allow right decision if left was in the last actions
+                if (lastActions.Contains(MovementState.Left))
+                {
+                    SetActionMask(RIGHT);
+                }
+            }
         }
 
         // to be implemented by the developer
@@ -129,25 +139,21 @@ namespace Cozmo
                 case STOP:
                     movementController.currentMovementState = MovementState.Stop;
                     lastChosenMovement = MovementState.Stop;
-                    //Test
                     SetReward(-0.02f);
                     break;
                 case FORWARD:
                     movementController.currentMovementState = MovementState.Forward;
                     lastChosenMovement = MovementState.Forward;
-                    //Test
                     SetReward(0.02f);
                     break;
                 case RIGHT:
                     movementController.currentMovementState = MovementState.Right;
                     lastChosenMovement = MovementState.Right;
-                    //Test
                     SetReward(0.01f);
                     break;
                 case LEFT:
                     movementController.currentMovementState = MovementState.Left;
                     lastChosenMovement = MovementState.Left;
-                    //Test
                     SetReward(0.01f);
                     break;
                 default:
@@ -155,7 +161,8 @@ namespace Cozmo
                     throw new ArgumentException("Invalid action value. Stop movement.");
             }
 
-            CollectLastMovementStates(lastChosenMovement);
+            if (!useOriginalActionMasking)
+                CollectLastMovementStates(lastChosenMovement);
 
             // Render new image after movement in order to update the centerOfGravity
             if (renderCamera != null)