static/demo.min.js

(() => {

  // Original gamepad logic: https://gist.github.com/videlais/8110000
  // Modified from: https://codepen.io/XanderLuciano/pen/vddOOG?editors=0010
  class GamePad {
    constructor() {
      this.supported = (navigator.webkitGetGamepads && navigator.webkitGetGamepads()) ||
        !!navigator.webkitGamepads || !!navigator.mozGamepads ||
        !!navigator.msGamepads || !!navigator.gamepads ||
        (navigator.getGamepads && navigator.getGamepads());

      this.ticking = false;
      this.connected = false;

      this.lStick = new THREE.Vector3(0,0,0);
      this.rStick = new THREE.Vector3(0,0,0);

      // Recommended deadzones for Xbox One controller
      this.SPACEMOUSE_THRESHOLD = 3000 / 32767.0;

      this.gamepads = [];
      this.prevRawGamepadTypes = [];
      this.prevTimestamps = [];

      this.init();
    }

    init() {
      if (this.supported) {
        // Older Firefox 
        window.addEventListener('MozGamepadConnected', (e) => this.onGamepadConnect(e), false);
        window.addEventListener('MozGamepadDisconnected', (e) => this.onGamepadDisconnect(e), false);

        //W3C Specification
        window.addEventListener('gamepadconnected', (e) => this.onGamepadConnect(e), false);
        window.addEventListener('gamepaddisconnected', (e) => this.onGamepadDisconnect(e), false);

        // Chrome
        if (navigator.webkitGetGamepads && navigator.webkitGetGamepads()) {
          this.startPolling();
        }

        //CocoonJS
        if(navigator.getGamepads && navigator.getGamepads()) {
          this.startPolling();
        }
      } else {
        console.log('Gamepad API not supported or not detected!');
      }
    }

    startPolling() {
      if (!this.ticking) {
        this.ticking = true;
        this.update();
      }
    }

    stopPolling() {
      this.ticking = false;
    }

    // Called externally
    update() {
      this.pollStatus();
      if (this.ticking) {
        this.pollJoysticks();
      }
    }

    pollStatus() {
      this.pollGamepads();
      for (let i in this.gamepads) {
        let gamepad = this.gamepads[i];
        if (gamepad.timestamp && (gamepad.timestamp === this.prevTimestamps[i])) {
          continue;
        }
        this.prevTimestamps[i] = gamepad.timestamp;
      }
    }

    pollGamepads() {
      let rawGamepads = (navigator.webkitGetGamepads && navigator.webkitGetGamepads()) ||
          navigator.webkitGamepads || navigator.mozGamepads ||
          navigator.msGamepads || navigator.gamepads || 
          (navigator.getGamepads && navigator.getGamepads());
      if (rawGamepads) {
        this.gamepads = [];
        for (let i = 0, max = rawGamepads.length; i < max; i++) {
          if (typeof rawGamepads[i] !== this.prevRawGamepadTypes[i]) {
            this.prevRawGamepadTypes[i] = typeof rawGamepads[i];
          }
          if (rawGamepads[i]) {
            this.gamepads.push(rawGamepads[i]);
          }
        }
      }
    }

    pollJoysticks() {
      let pad = 0;

      // Reset all input to 0
      // this.lStick.set(0, 0, 0);
      this.rStick.set(0, 0, 0);

      if (this.gamepads[pad]) {

        let panX  = this.gamepads[pad].axes[0];
        let panY  = this.gamepads[pad].axes[1];
        let rollX = this.gamepads[pad].axes[2];
        let rollY = this.gamepads[pad].axes[3];

        let buttons = this.gamepads[pad].buttons;
        this.rStick.z = buttons[11].pressed ? 1 : 0;

        if (panX < -this.SPACEMOUSE_THRESHOLD ||
            panX > this.SPACEMOUSE_THRESHOLD) {
          this.lStick.x = panX;
        }

        if (panY < -this.SPACEMOUSE_THRESHOLD ||
            panY > this.SPACEMOUSE_THRESHOLD) {
          this.lStick.y = -panY;
        }

        if (rollX < -this.SPACEMOUSE_THRESHOLD ||
            rollX > this.SPACEMOUSE_THRESHOLD) {
          this.rStick.x = rollX;
        }

        if (rollY < -this.SPACEMOUSE_THRESHOLD ||
            rollY > this.SPACEMOUSE_THRESHOLD) {
          this.rStick.y = -rollY;
        }
      }
    }

    onGamepadConnect(event) {
      console.log('Gamepad Connected!');
      this.connected = true;
      let gamepad = event.gamepad;
      this.gamepads[event.gamepad.id] = gamepad;
      this.startPolling();
    }

    onGamepadDisconnect(event) {
      console.log('Gamepad Disconnected!');
      this.connected = false;
      this.gamepads[event.gamepad.id] = null;
      if (this.gamepads.length === 0) {
        this.stopPolling();
      }
    }
  }

  const button = document.getElementById("demoButton");
  button.onclick = async function (event) {
    // Prevent loading multiple instances
    button.onclick = null;
    button.innerText = "Loading Modules..."
    setTimeout(() => {button.innerText = "If demo is not loading, visit project website link below"}, 10000)

    await Promise.all([
      import("./onnx.min.js"),
      import("./three/build/three.module.min.js"),
      import("./three/examples/jsm/controls/OrbitControls.min.js"),
      import("./three/examples/jsm/libs/dat.gui.module.min.js"),
    ]).then((results) => {
      window.THREE = results[1];
      window.OrbitControls = results[2].OrbitControls;
      window.GUI = results[3].GUI;
    });

    const modelFiles = await Promise.all([
      fetch("https://huggingface.co/spaces/belinghy/character-animation-motion-vaes/resolve/main/static/mvae.onnx").then((f) => f.arrayBuffer()),
      fetch("https://huggingface.co/spaces/belinghy/character-animation-motion-vaes/resolve/main/static/target_controller.onnx").then((f) => f.arrayBuffer()),
      fetch("https://huggingface.co/spaces/belinghy/character-animation-motion-vaes/resolve/main/static/joystick_controller.onnx").then((f) => f.arrayBuffer()),
    ]).then((files) => {
      return {
        mvae: files[0],
        target: files[1],
        joystick: files[2]
      }
    })

    // Create Onnx model
    const onnxSessions = {};
    for (let [key, file] of Object.entries(modelFiles)) {
      onnxSessions[key] = new onnx.InferenceSession({
        backendHint: "webgl",
      });
      await onnxSessions[key].loadModel(file);
    }

    // Connect Controller
    const gamepad = new GamePad();
    window.gamepad = gamepad;

    // Make ThreeJS scene
    const container = document.getElementById("mvaeScene")
    const renderer = new THREE.WebGLRenderer({
      antialias: true,
      alpha: true,
    });
    renderer.setPixelRatio(window.devicePixelRatio);
    renderer.setSize(container.clientWidth, container.clientHeight);
    renderer.outputEncoding = THREE.sRGBEncoding;
    renderer.shadowMap.enabled = true;
    renderer.domElement.style.outline = "none";
    container.appendChild(renderer.domElement);

    const scene = new THREE.Scene();
    scene.background = new THREE.Color(0x282c34);
    scene.fog = new THREE.Fog(0x282c34, 10, 400);

    let aspectRatio = container.clientWidth / container.clientHeight;
    const camera = new THREE.PerspectiveCamera(25, aspectRatio, 0.1, 1000);
    camera.lookAt(0, 0, 0);
    camera.position.set(-30, 10, 0);
    const controls = new OrbitControls(camera, renderer.domElement);
    controls.enableDamping = true;
    controls.minDistance = 20;
    controls.maxDistance = 200;
    controls.update();

    // Utilities
    const loader = new THREE.TextureLoader();
    const raycaster = new THREE.Raycaster();
    const mouse = new THREE.Vector2();

    // Character states
    let speed = 1;
    let heading = 0;
    let rootFacing = 0;
    let rootXYZ = [0, 0, 0];
    let action = new onnx.Tensor(new Float32Array(32), "float32", [1, 32]);
    let condition = new onnx.Tensor(new Float32Array(267), "float32", [1, 267]);

    const poseStd = [0.2337,0.0663,0.0614,0.9697,0.2778,0.7038,0.9697,0.2778,0.7038,0.1865,0.3471,0.1654,0.2632,0.3683,0.2215,0.3824,0.408,0.3032,0.3247,0.3808,0.2617,0.3247,0.3808,0.2617,0.3628,0.3749,0.2736,0.3628,0.3749,0.2736,0.4146,0.3826,0.2113,0.4146,0.3826,0.2113,0.0048,0.3071,0.0048,0.7615,0.3368,0.6285,0.7615,0.3368,0.6285,0.0158,0.3156,0.016,0.0158,0.3156,0.016,0.4892,0.2069,0.3532,0.4892,0.2069,0.3532,0.0492,0.3164,0.0526,0.12,0.3314,0.1132,0.5161,0.4031,0.2726,0.5161,0.4031,0.2726,0.4056,0.1132,0.1208,0.4056,0.1132,0.1208,0.2268,0.054,0.0672,0.2245,0.0548,0.0669,0.2213,0.0567,0.0698,0.223,0.0554,0.0677,0.223,0.0554,0.0677,0.2255,0.0607,0.0691,0.2255,0.0607,0.0691,0.2488,0.1168,0.0937,0.2488,0.1168,0.0937,0.2337,0.0527,0.0714,0.3489,0.1147,0.1136,0.3489,0.1147,0.1136,0.2348,0.0533,0.0728,0.2348,0.0533,0.0728,0.291,0.0778,0.081,0.291,0.0778,0.081,0.2317,0.0529,0.0705,0.229,0.0534,0.0686,0.2855,0.1421,0.1301,0.2855,0.1421,0.1301,0.4053,0.3292,0.1624,0.4053,0.3292,0.1624,0.0919,0.2076,0.0945,0.1048,0.2116,0.1123,0.1294,0.1692,0.1227,0.1161,0.2072,0.1589,0.1161,0.2072,0.1589,0.2761,0.2619,0.2591,0.2761,0.2619,0.2591,0.397,0.3733,0.2129,0.397,0.3733,0.2129,0.0409,0.1883,0.0369,0.243,0.4652,0.1788,0.243,0.4652,0.1788,0.1376,0.3043,0.2111,0.1376,0.3043,0.2111,0.2872,0.3043,0.2376,0.2872,0.3043,0.2376,0.0686,0.2018,0.0612,0.0799,0.2049,0.0769,0.4131,0.4178,0.326,0.413,0.4178,0.326,0.1946,0.2185,0.0824,0.1946,0.2185,0.0824,0.083,0.1624,0.0254,0.0995,0.1588,0.0252,0.1439,0.1504,0.0314,0.2075,0.1172,0.1885,0.2075,0.1172,0.1885,0.2097,0.2119,0.1359,0.2097,0.2119,0.1359,0.2097,0.2119,0.1359,0.2097,0.2119,0.1359,0.0372,0.2014,0.0372,0.1857,0.2257,0.0808,0.1857,0.2257,0.0808,0.1552,0.2805,0.094,0.1552,0.2805,0.094,0.1552,0.2805,0.094,0.1552,0.2805,0.094,0.052,0.1751,0.0286,0.0664,0.1678,0.0266,0.2818,0.2703,0.1431,0.2818,0.2703,0.1431];
    const poseAvg = [2.6107e-1,-4.7486e-7,1.401e-7,-2.8028e-3,3.0103e-1,-3.6838e-1,-2.7801e-3,3.0103e-1,3.6839e-1,1.4925e-1,3.8383,-8.6033e-7,2.4685e-1,4.1773,-6.7364e-7,4.1594e-1,4.6868,-2.4816e-8,2.8113e-1,4.435,-8.8858e-2,2.8113e-1,4.435,8.8857e-2,2.9221e-1,4.4636,-5.9573e-1,2.9222e-1,4.4636,5.9572e-1,-1.1323e-1,3.8501,9.0934e-1,-1.1324e-1,3.8501,-9.0934e-1,-5.509e-5,2.9514,7.9133e-12,-3.1353e-1,6.6536e-1,-2.7751e-1,-3.1351e-1,6.6536e-1,2.7751e-1,-4.9933e-2,2.9565,-3.931e-1,-4.9933e-2,2.9565,3.931e-1,3.8439e-1,1.6758,-4.3913e-1,3.844e-1,1.6758,4.3913e-1,8.3124e-3,3.2001,-5.0445e-7,7.5637e-2,3.536,-8.0779e-7,3.6293e-1,3.3118,-7.784e-1,3.6294e-1,3.3118,7.784e-1,2.8288e-1,1.953e-4,-1.3949e-2,2.8288e-1,1.9532e-4,1.3949e-2,2.5435e-1,-1.1878e-4,-5.0363e-7,2.5261e-1,-1.0782e-4,-5.3367e-7,2.5019e-1,-8.4664e-5,-5.927e-7,2.5142e-1,-1.101e-4,8.511e-5,2.5142e-1,-1.1008e-4,-8.6227e-5,2.5097e-1,-5.3605e-5,1.2627e-3,2.5097e-1,-5.3493e-5,-1.2638e-3,2.5386e-1,2.9768e-5,-5.1661e-3,2.5386e-1,2.9656e-5,5.1651e-3,2.5973e-1,-1.2142e-4,-4.7796e-7,2.802e-1,2.5689e-4,-1.1458e-2,2.802e-1,2.5717e-4,1.1458e-2,2.5962e-1,-1.3255e-4,-5.8587e-4,2.5962e-1,-1.3255e-4,5.8493e-4,2.7278e-1,1.8581e-5,-9.093e-3,2.7278e-1,1.8695e-5,9.0928e-3,2.5799e-1,-1.2785e-4,-4.7312e-7,2.5601e-1,-1.2486e-4,-4.8496e-7,2.5366e-1,6.089e-4,7.9774e-3,2.5366e-1,6.0939e-4,-7.9785e-3,-2.0331e-1,-8.1218e-1,7.5829e-4,-2.0331e-1,-8.1218e-1,-7.5517e-4,9.3253e-1,-2.6446e-1,-1.7862e-6,9.1736e-1,-3.0011e-1,-2.2849e-6,8.3715e-1,-4.886e-1,-9.844e-6,9.0167e-1,-3.2442e-1,-7.9873e-3,9.0167e-1,-3.2442e-1,7.9802e-3,6.6979e-1,-5.487e-1,1.958e-1,6.6978e-1,-5.487e-1,-1.9581e-1,5.1202e-1,5.3974e-1,-3.2292e-1,5.1202e-1,5.3974e-1,3.2291e-1,9.7553e-1,-9.9292e-2,-4.7701e-8,-6.1732e-1,-5.5017e-1,9.3872e-2,-6.1732e-1,-5.5018e-1,-9.3873e-2,-8.5457e-1,-2.7279e-1,1.9796e-1,-8.5457e-1,-2.7279e-1,-1.9796e-1,-7.157e-1,4.8543e-1,1.4351e-1,-7.1571e-1,4.8542e-1,-1.4351e-1,9.5673e-1,-1.8836e-1,-8.6077e-7,9.4572e-1,-2.2657e-1,-1.3113e-6,4.6985e-1,5.1673e-1,2.4654e-1,4.6985e-1,5.1673e-1,-2.4655e-1,-1.6761e-1,3.9161e-2,-9.3701e-1,1.6762e-1,-3.9163e-2,-9.3701e-1,-1.7539e-6,6.0989e-7,-9.8291e-1,-2.1627e-6,1.3742e-6,-9.8196e-1,-1.0023e-5,1.1619e-6,-9.7758e-1,-3.2096e-1,-8.9459e-1,-6.6042e-2,3.2096e-1,8.9459e-1,-6.6042e-2,-3.7467e-1,-1.8443e-1,8.4751e-1,3.7468e-1,1.8443e-1,8.475e-1,3.7468e-1,1.8443e-1,8.475e-1,-3.7467e-1,-1.8443e-1,8.4751e-1,4.7733e-8,-1.9534e-6,-9.781e-1,-1.6384e-1,4.3208e-2,-9.3774e-1,1.6385e-1,-4.321e-2,-9.3774e-1,-1.9273e-1,-2.0136e-2,-9.224e-1,1.9273e-1,2.0141e-2,-9.224e-1,-1.9272e-1,-2.0145e-2,-9.224e-1,1.9273e-1,2.015e-2,-9.224e-1,-8.4422e-7,-8.6653e-7,-9.8276e-1,-1.3142e-6,-1.3712e-7,-9.8322e-1,-2.5745e-1,-2.0945e-1,8.467e-1,2.5746e-1,2.0945e-1,8.4669e-1];
    const startPose = [3.7343e-1,2.5392e-2,-3.6897e-3,9.1467e-1,2.0899e-1,-1.0637,-1.6662,7.3143e-1,-8.8882e-1,2.9283e-1,3.852,2.7816e-1,4.5728e-1,4.1711,3.5562e-1,7.3643e-1,4.6469,4.5095e-1,5.4106e-1,4.4424,3.2156e-1,5.5217e-1,4.4007,4.9731e-1,5.3249e-1,4.5413,-1.9381e-1,6.4737e-1,4.3115,1.0057,8.8589e-2,3.682,1.3214,2.2347e-1,4.0642,-8.903e-1,0.0,3.0089,0.0,4.1466e-1,5.0355e-1,-9.374e-1,-1.4528,1.225,-6.3671e-1,-3.1384e-2,3.1669,-3.7173e-1,-6.7839e-2,2.8752,3.7637e-1,6.6035e-1,1.8958,-7.6142e-1,-1.062e-1,1.5319,-2.8707e-1,4.4212e-2,3.2461,9.468e-2,1.651e-1,3.5656,1.9719e-1,7.4468e-1,3.3334,-8.4051e-1,6.4904e-1,3.2906,7.3755e-1,1.7409e-1,-1.2108e-1,2.5089e-2,3.8499e-1,2.8025e-1,1.2745e-1,3.6488e-1,-1.5865e-2,-3.3239e-3,3.6296e-1,-1.1497e-2,-1.6055e-2,3.6115e-1,-5.9476e-3,-3.615e-2,3.6703e-1,-1.4265e-2,-3.2338e-2,3.5635e-1,-9.9516e-4,-2.7624e-2,4.116e-1,-6.9391e-2,-4.8794e-2,3.0963e-1,2.5622e-2,-1.1095e-2,2.3451e-1,9.2727e-2,6.2983e-3,5.5357e-1,-2.2211e-1,-3.8469e-2,3.7352e-1,-2.8061e-2,2.4014e-2,2.0644e-1,-7.8941e-2,7.6509e-2,4.9427e-1,2.346e-1,1.5348e-1,3.7234e-1,-3.8626e-2,1.9442e-2,3.7435e-1,-1.8714e-2,2.7618e-2,2.8971e-1,-7.6181e-2,-1.701e-2,6.4407e-1,-3.4401e-2,-1.2438e-2,3.7001e-1,-2.4745e-2,1.7678e-2,3.6711e-1,-2.0063e-2,7.1973e-3,7.4105e-1,-1.1856e-1,3.7614e-2,1.6907e-1,-2.2738e-1,6.3841e-2,8.9425e-3,-9.9831e-1,-5.7434e-2,-9.8963e-1,-1.1627e-1,-8.4324e-2,8.9335e-1,-4.2454e-1,-1.4729e-1,8.646e-1,-4.7384e-1,-1.6715e-1,7.5405e-1,-6.4569e-1,-1.2038e-1,8.5399e-1,-5.0816e-1,-1.117e-1,8.4038e-1,-4.8479e-1,-2.4236e-1,6.5666e-1,-7.256e-1,2.0567e-1,4.5273e-1,-6.8694e-1,-5.6846e-1,4.5016e-1,8.79e-1,-1.572e-1,4.6176e-1,3.8382e-1,7.9967e-1,9.7032e-1,-2.3754e-1,-4.5353e-2,-4.9521e-1,-8.6664e-1,6.0737e-2,-9.3292e-1,3.3144e-1,1.4074e-1,-8.5423e-1,-5.0399e-1,1.2763e-1,-9.9094e-1,8.114e-2,-1.0699e-1,-9.5551e-1,1.355e-1,2.6199e-1,-2.9873e-1,8.7383e-1,3.8365e-1,9.4078e-1,-3.2382e-1,-1.0032e-1,9.1878e-1,-3.7446e-1,-1.2497e-1,4.6267e-1,5.1188e-1,7.2382e-1,6.0684e-1,7.2977e-1,3.1492e-1,-2.146e-1,5.4182e-2,-9.752e-1,2.3719e-2,4.4674e-1,-8.9435e-1,-3.8452e-2,2.5434e-1,-9.6635e-1,-6.1361e-2,2.306e-1,-9.7111e-1,8.8725e-3,1.9328e-1,-9.811e-1,-5.2004e-1,-8.404e-1,-1.5259e-1,5.1074e-1,8.5799e-1,5.4759e-2,-6.713e-1,-4.3805e-1,5.9789e-1,6.3929e-1,-1.9434e-1,7.44e-1,6.3929e-1,-1.9434e-1,7.44e-1,-6.713e-1,-4.3805e-1,5.9789e-1,4.5352e-2,3.6296e-1,-9.307e-1,-2.144e-1,5.4166e-2,-9.7524e-1,2.3728e-2,4.4659e-1,-8.9442e-1,-2.3929e-1,1.6321e-1,-9.5713e-1,1.3181e-1,4.3593e-1,-8.9027e-1,-2.3929e-1,1.6321e-1,-9.5713e-1,1.3181e-1,4.3593e-1,-8.9027e-1,3.2209e-3,3.0445e-1,-9.5252e-1,-1.6893e-2,2.79e-1,-9.6014e-1,-3.5985e-1,-6.3774e-1,6.8102e-1,3.5872e-1,-6.0504e-1,7.1081e-1];
    const linkedJoints = [
      [12, 0], // right foot
      [16, 12], // right shin
      [14, 16], // right leg
      [15, 17], // left foot
      [17, 13], // left shin
      [13, 1], // left leg
      [5, 7], // right shoulder
      [7, 10], // right upper arm
      [10, 20], // right lower arm
      [6, 8], // left shoulder
      [8, 9], // left upper arm
      [9, 21], // left lower arm
      [3, 18], // torso
      [14, 15], // hip
    ];
    const F2M = 0.3048;
    const M2F = 1 / 0.3048;
    const TAU = 2 * Math.PI;
    const GREY = 0x607d8b;

    function resetPose() {
      for (let i = 0; i < condition.data.length; i++) {
        condition.set((startPose[i] - poseAvg[i]) / poseStd[i], [0, i]);
      }
    }

    function resetState() {
      rootFacing = 0;
      rootXYZ = [0, 0, 0];
      controls.target.x = 0;
      controls.target.z = 0;
      target.position.x = 15;
      target.position.z = 15;
      speed = 1;
      heading = 0;
      resetPose();

      if (guiParams.Task == "target") {
        camera.position.set(30, 10, 30);
      } else {
        camera.position.set(-30, 10, 0);
      }
    }

    // Make GUI before scene
    const gui = new GUI({ autoPlace: false });
    let guiParams = {
      Task: "target",
      Reset: false,
      Pause: false,
    };
    gui
      .add(guiParams, "Task", {
        Target: "target",
        Joystick: "joystick",
      })
      .setValue("joystick")
      .onChange(() => {
        resetState();
        renderer.domElement.focus();
      });

    const resetCon = gui.add(guiParams, "Reset");
    resetCon.onChange(() => {
      resetState();
      setTimeout(() => {
        guiParams.Reset = false;
        resetCon.updateDisplay();
      }, 100);
    });
    gui.add(guiParams, "Pause");
    document.getElementById("guiContainer").appendChild(gui.domElement);

    // SpotLight
    const spotLight = new THREE.SpotLight(0xffffff);
    spotLight.position.set(100, 1000, 100);
    spotLight.castShadow = true;
    scene.add(spotLight);

    // Ground
    const groundTexture = loader.load("./static/grid.png");
    groundTexture.wrapS = groundTexture.wrapT = THREE.RepeatWrapping;
    groundTexture.repeat.set(50, 50);
    groundTexture.anisotropy = 16;
    groundTexture.encoding = THREE.sRGBEncoding;
    const ground = new THREE.Mesh(
      new THREE.PlaneBufferGeometry(400, 400),
      new THREE.MeshLambertMaterial({ map: groundTexture })
    );
    ground.rotation.x = -Math.PI / 2;
    ground.receiveShadow = true;
    scene.add(ground);

    // Target
    const target = new THREE.Mesh(
      new THREE.SphereBufferGeometry(0.15, 16, 16),
      new THREE.MeshBasicMaterial({ color: 0xff0000 })
    );
    scene.add(target);

    // Character
    const jointGeom = new THREE.SphereBufferGeometry(0.07, 16, 16);
    const jointMat = new THREE.MeshBasicMaterial({ color: 0x343837 });
    const joints = [];
    for (let i = 0; i < 22; i++) {
      const sphere = new THREE.Mesh(jointGeom, jointMat);
      scene.add(sphere);
      joints.push(sphere);
    }

    const linkGeom = new THREE.CylinderBufferGeometry(0.06, 0.06, 1, 16);
    const linkMat = new THREE.MeshBasicMaterial({ color: 0x3e82fc });
    const links = [];
    for (let i = 0; i < linkedJoints.length; i++) {
      const cylinder = new THREE.Mesh(linkGeom, linkMat);
      scene.add(cylinder);
      links.push(cylinder);
    }

    const head = new THREE.Mesh(
      new THREE.SphereBufferGeometry(0.125, 16, 16),
      new THREE.MeshBasicMaterial({ color: 0x3e82fc })
    );
    scene.add(head);

    resetState();

    // Event Handlers
    function setTarget(eventX, eventY) {
      const canvas = renderer.domElement;
      const bbox = canvas.getBoundingClientRect();
      mouse.x = ((eventX - bbox.x) / canvas.clientWidth) * 2 - 1;
      mouse.y = -((eventY - bbox.y) / canvas.clientHeight) * 2 + 1;

      raycaster.setFromCamera(mouse, camera);
      const intersects = raycaster.intersectObjects([ground]);
      for (const intersect of intersects) {
        target.position.x = intersect.point.x;
        target.position.z = intersect.point.z;
      }
      if (guiParams.Task == "joystick") {
        heading = Math.atan2(
          target.position.z - rootXYZ[2],
          target.position.x - rootXYZ[0]
        );
      }
    }

    function onWindowResize() {
      camera.aspect = container.clientWidth / container.clientHeight;
      camera.updateProjectionMatrix();
      renderer.setSize(container.clientWidth, container.clientHeight);
    }
    function onMouseClick(event) {
      event.preventDefault();
      const x = event.clientX;
      const y = event.clientY;
      setTarget(x, y);
    }
    function onTouchStart(event) {
      event.preventDefault();
      const x = event.touches[0].clientX;
      const y = event.touches[0].clientY;
      setTarget(x, y);
    }

    // Start animation loop
    window.addEventListener("resize", onWindowResize, false);
    renderer.domElement.addEventListener("auxclick", onMouseClick, false);
    renderer.domElement.addEventListener("touchstart", onTouchStart, false);

    // 30 fps in milliseconds
    const fpsInterval = 1000 / 30;
    let startTime = performance.now();

    const animate = async () => {
      requestAnimationFrame(animate);

      let endTime = performance.now();
      let elapsedTime = endTime - startTime;

      if (!guiParams.Pause && elapsedTime > fpsInterval) {
        startTime = endTime - (elapsedTime % fpsInterval);

        // Get rotation matrix [a b; c d]
        const a = Math.cos(rootFacing);
        const b = Math.sin(rootFacing);
        const c = -b;
        const d = a;

        // First calculate observation components
        let taskVector = [0.0, 0.0];
        if (guiParams.Task == "target") {
          let targetDelta = [
            target.position.x - rootXYZ[0],
            target.position.z - rootXYZ[2],
          ];
          let distance = Math.sqrt(targetDelta.reduce((a, c) => a + c * c, 0));
          if (distance * M2F < 2.0) {
            do {
              target.position.x = (Math.random() * 80.0 - 40.0) * F2M;
              target.position.z = (Math.random() * 80.0 - 40.0) * F2M;
              targetDelta = [
                target.position.x - rootXYZ[0],
                target.position.z - rootXYZ[2],
              ];
              distance = Math.sqrt(targetDelta.reduce((a, c) => a + c * c, 0));
            } while (distance * M2F < 20.0);
          }
          taskVector[0] = (a * targetDelta[0] + c * targetDelta[1]) * M2F;
          taskVector[1] = (b * targetDelta[0] + d * targetDelta[1]) * M2F;
        } else if (guiParams.Task == "joystick") {

          let dz = camera.position.z - controls.target.z;
          let dx = camera.position.x - controls.target.x;
          let cameraAngle = Math.atan2(dz, dx);

          if (gamepad.connected) { gamepad.update(); }

          heading = Math.atan2(gamepad.lStick.x, gamepad.lStick.y) + (cameraAngle - Math.PI);
          speed = Math.min(Math.max(gamepad.lStick.length(), 0.2), 0.8);
          cameraAngle += gamepad.rStick.x * 5 / 360 * Math.PI;
          if (gamepad.rStick.z == 1) { cameraAngle = Math.PI; }

          target.position.x = rootXYZ[0] + (2 * speed + 0.6) * Math.cos(heading);
          target.position.z = rootXYZ[2] + (2 * speed + 0.6) * Math.sin(heading);
          const targetAngle = heading + rootFacing;
          speed = Math.floor(speed * 10) / 10; // task vector needs discrete speeds
          taskVector[0] = speed * Math.cos(targetAngle);
          taskVector[1] = speed * Math.sin(targetAngle);

          let cx = rootXYZ[0] + Math.cos(cameraAngle) * 30;
          let cz = rootXYZ[2] + Math.sin(cameraAngle) * 30; 
          camera.position.set(cx, 10, cz);

          controls.target.x = rootXYZ[0];
          controls.target.z = rootXYZ[2];
        } else {
          console.log(`Oops, "${guiParams.Task}" if not a supported task.`);
        }

        // Get action from policy
        let controllerInput = new onnx.Tensor(new Float32Array(269), "float32", [1, 269]);
        for (let i = 0; i < condition.size; i++) {
          // Controller input is denormalized
          const c = condition.get([0, i]) * poseStd[i] + poseAvg[i];
          controllerInput.set(c, [0, i]);
        }
        for (let i = 0; i < taskVector.length; i++) {
          const c = taskVector[i] + 0.00001; // Make sure it's float
          controllerInput.set(c, [0, condition.size + i]);
        }
        await onnxSessions[guiParams.Task]
          .run([controllerInput])
          .then((output) => {
            action = output.values().next().value;
          });

        // Get next frame from decoder
        await onnxSessions.mvae.run([action, condition]).then((output) => {
          condition = output.values().next().value;
        });

        const pose = condition.data;

        // Integrate root translation
        // Denormalize and convert to meters for joint positions
        const normPose = pose.slice(0, 69).map((p, i) => {
          return (p * poseStd[i] + poseAvg[i]) * F2M;
        });
        // "Unconvert" root facing delta
        normPose[2] = normPose[2] / F2M;

        const gdx = normPose[0];
        const gdz = normPose[1];
        const ldx = a * gdx + b * gdz; // local frame
        const ldz = c * gdx + d * gdz; // local frame
        rootXYZ[0] += ldx;
        rootXYZ[2] += ldz;
        rootFacing = (rootFacing + normPose[2]) % TAU;

        // Set joint positions
        for (let i = 0; i < joints.length; i++) {
          const j = (i + 1) * 3;
          const k = j + 1;
          const l = j + 2;

          // Rotate joint positions to global frame
          const gx = normPose[j];
          const gz = normPose[l];
          const lx = a * gx + b * gz;
          const lz = c * gx + d * gz;

          // Render needs denormalized conditions
          joints[i].position.x = lx + rootXYZ[0];
          joints[i].position.y = normPose[k];
          joints[i].position.z = lz + rootXYZ[2];
        }

        // Set joint links
        for (let i = 0; i < links.length; i++) {
          const [j, k] = linkedJoints[i];
          const [jx, jy, jz] = [
            joints[j].position.x,
            joints[j].position.y,
            joints[j].position.z,
          ];
          const [kx, ky, kz] = [
            joints[k].position.x,
            joints[k].position.y,
            joints[k].position.z,
          ];
          const position = [(kx + jx) / 2, (ky + jy) / 2, (kz + jz) / 2];
          const delta = [kx - jx, ky - jy, kz - jz];
          const length = Math.sqrt(delta.reduce((a, c) => a + c * c, 0));
          const quaternion = [delta[2], 0, -delta[0], length + delta[1]];
          const qnorm = Math.sqrt(quaternion.reduce((a, c) => a + c * c, 0));

          links[i].scale.y = length;
          links[i].position.x = position[0];
          links[i].position.y = position[1];
          links[i].position.z = position[2];
          links[i].quaternion.x = quaternion[0] / qnorm;
          links[i].quaternion.y = quaternion[1] / qnorm;
          links[i].quaternion.z = quaternion[2] / qnorm;
          links[i].quaternion.w = quaternion[3] / qnorm;
        }

        // Set head position
        head.position.x = 1.75 * joints[4].position.x - 0.75 * joints[3].position.x;
        head.position.y = 1.75 * joints[4].position.y - 0.75 * joints[3].position.y;
        head.position.z = 1.75 * joints[4].position.z - 0.75 * joints[3].position.z;

        // must be called after any manual changes to the camera
        // or if enableDamping is true
        controls.update();
        renderer.render(scene, camera);
      }
    };

    button.remove();
    animate();
  }
})();