update pushblock actions

Project/Assets/ML-Agents/Examples/PushBlock/Scripts/PushAgentCollab.cs

         var dirToGo = Vector3.zero;
         var rotateDir = Vector3.zero;
 
-        var action = act[0];
+
+        var forwardAxis = act[0];
+        var rightAxis = act[1];
+        var rotateAxis = act[2];
-        switch (action)
+        switch (forwardAxis)
         {
             case 1:
                 dirToGo = transform.forward * 1f;
                 break;
-            case 3:
-                rotateDir = transform.up * 1f;
+        }
+
+        switch (rightAxis)
+        {
+            case 1:
+                dirToGo = transform.right * 1f;
+                break;
+            case 2:
+                dirToGo = transform.right * -1f;
-            case 4:
+        }
+
+        switch (rotateAxis)
+        {
+            case 1:
-            case 5:
-                dirToGo = transform.right * -0.75f;
-                break;
-            case 6:
-                dirToGo = transform.right * 0.75f;
+            case 2:
+                rotateDir = transform.up * 1f;
+
         transform.Rotate(rotateDir, Time.fixedDeltaTime * 200f);
         m_AgentRb.AddForce(dirToGo * m_PushBlockSettings.agentRunSpeed,
             ForceMode.VelocityChange);
     public override void Heuristic(in ActionBuffers actionsOut)
     {
         var discreteActionsOut = actionsOut.DiscreteActions;
-        discreteActionsOut[0] = 0;
-        if (Input.GetKey(KeyCode.D))
+        discreteActionsOut.Clear();
+        //forward
+        if (Input.GetKey(KeyCode.W))
+        {
+            discreteActionsOut[0] = 1;
+        }
+        if (Input.GetKey(KeyCode.S))
+        {
+            discreteActionsOut[0] = 2;
+        }
+        //rotate
+        if (Input.GetKey(KeyCode.A))
-            discreteActionsOut[0] = 3;
+            discreteActionsOut[2] = 1;
-        else if (Input.GetKey(KeyCode.W))
+        if (Input.GetKey(KeyCode.D))
-            discreteActionsOut[0] = 1;
+            discreteActionsOut[2] = 2;
-        else if (Input.GetKey(KeyCode.A))
+        //right
+        if (Input.GetKey(KeyCode.E))
-            discreteActionsOut[0] = 4;
+            discreteActionsOut[1] = 1;
-        else if (Input.GetKey(KeyCode.S))
+        if (Input.GetKey(KeyCode.Q))
-            discreteActionsOut[0] = 2;
+            discreteActionsOut[1] = 2;
         }
     }