Fix a UI bug and beautify cart-pole rendering

cart-pole/index.html

@@ -55,6 +55,9 @@
       display: inline-block;
       padding-left: 0px;
       padding-right: 10px;
+      padding-top: 10px;
+      padding-bottom: 10px;
+      margin: 5px;
       vertical-align: top;
       border: 1px #AAA solid;
     }

cart-pole/index.js

@@ -134,13 +134,10 @@ class PolicyNetwork {
 
         await maybeRenderDuringTraining(cartPoleSystem);
 
-        if (isDone || j >= maxStepsPerGame) {
-          if (isDone) {
-            // When the game ends before max step count is reached, a reward of
-            // 0 is given.
-            gameRewards.push(0);
-          }
-          onGameEnd(i + 1, numGames);
+        if (isDone) {
+          // When the game ends before max step count is reached, a reward of
+          // 0 is given.
+          gameRewards.push(0);
           break;
         } else {
           // As long as the game doesn't end, each step leads to a reward of 1.
@@ -149,6 +146,7 @@ class PolicyNetwork {
           gameRewards.push(1);
         }
       }
+      onGameEnd(i + 1, numGames);
       gameSteps.push(gameRewards.length);
       this.pushGradients(allGradients, gameGradients);
       allRewards.push(gameRewards);

cart-pole/ui.js

@@ -145,27 +145,75 @@ function renderCartPole(cartPole, canvas) {
   context.clearRect(0, 0, canvas.width, canvas.height);
   const halfW = canvas.width / 2;
 
-  // 1. Draw the cart.
+  // Draw the cart.
   const railY = canvas.height * 0.8;
   const cartW = cartPole.cartWidth * scale;
   const cartH = cartPole.cartHeight * scale;
 
   const cartX = cartPole.x * scale + halfW;
 
   context.beginPath();
+  context.strokeStyle = '#000000';
+  context.lineWidth = 2;
   context.rect(cartX - cartW / 2, railY - cartH / 2, cartW, cartH);
   context.stroke();
 
-  // 2. Draw the pole.
+  // Draw the wheels under the cart.
+  const wheelRadius = cartH / 4;
+  for (const offsetX of [-1, 1]) {
+    context.beginPath();
+    context.lineWidth = 2;
+    context.arc(cartX - cartW / 4 * offsetX, railY + cartH / 2 + wheelRadius,
+                wheelRadius, 0, 2 * Math.PI);
+    context.stroke();
+  }
+
+  // Draw the pole.
   const angle = cartPole.theta + Math.PI / 2;
   const poleTopX =
       halfW + scale * (cartPole.x + Math.cos(angle) * cartPole.length);
   const poleTopY = railY -
       scale * (cartPole.cartHeight / 2 + Math.sin(angle) * cartPole.length);
   context.beginPath();
+  context.strokeStyle = '#ffa500';
+  context.lineWidth = 6;
   context.moveTo(cartX, railY - cartH / 2);
   context.lineTo(poleTopX, poleTopY);
   context.stroke();
+
+  // Draw the ground.
+  const groundY = railY + cartH / 2 + wheelRadius * 2;
+  context.beginPath();
+  context.strokeStyle = '#000000';
+  context.lineWidth = 1;
+  context.moveTo(0, groundY);
+  context.lineTo(canvas.width, groundY);
+  context.stroke();
+
+  const nDivisions = 40;
+  for (let i = 0; i < nDivisions; ++i) {
+    const x0 = canvas.width / nDivisions * i;
+    const x1 = x0 + canvas.width / nDivisions / 2;
+    const y0 = groundY + canvas.width / nDivisions / 2;
+    const y1 = groundY;
+    context.beginPath();
+    context.moveTo(x0, y0);
+    context.lineTo(x1, y1);
+    context.stroke();
+  }
+
+  // Draw the left and right limits.
+  const limitTopY = groundY - canvas.height / 2;
+  context.beginPath();
+  context.strokeStyle = '#ff0000';
+  context.lineWidth = 2;
+  context.moveTo(1, groundY);
+  context.lineTo(1, limitTopY);
+  context.stroke();
+  context.beginPath();
+  context.moveTo(canvas.width - 1, groundY);
+  context.lineTo(canvas.width - 1, limitTopY);
+  context.stroke();
 }
 
 async function updateUIControlState() {