waymax_demo_train.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/9Tempest/motionLM-Serve/blob/main/waymax_demo_train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rNYRA6k8Qfyo"
      },
      "source": [
        "# Scenario Data Loading\n",
        "\n",
        "This tutorial demonstrates how to load scenario data from the Waymo Open Motion Dataset (WOMD) using the Waymax dataloader."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8w2UKUVC4rs6"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/9Tempest/motionLM-Serve/blob/main/waymax_demo.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/colab_logo_1x.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MtgRcYqmtMwD"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install mediapy\n",
        "!pip install git+https://github.com/waymo-research/waymax.git@main#egg=waymo-waymax\n",
        "import numpy as np\n",
        "import mediapy\n",
        "from tqdm import tqdm\n",
        "import dataclasses\n",
        "import jax\n",
        "from jax import numpy as jnp\n",
        "import numpy as np\n",
        "import mediapy\n",
        "\n",
        "\n",
        "from waymax import config as _config\n",
        "from waymax import dataloader\n",
        "from waymax import datatypes\n",
        "from waymax import dynamics\n",
        "from waymax import env as _env\n",
        "from waymax import agents\n",
        "from waymax import visualization\n",
        "from google.colab import auth\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0o2sAapxRMAT"
      },
      "source": [
        "\n",
        "We first create a dataset config, using the default configs provided in the `waymax.config` module. In particular, `config.WOD_1_1_0_TRAINING` is a pre-defined configuration that points to version 1.1.0 of the Waymo Open Dataset.\n",
        "\n",
        "The data config contains a number of options to configure how and where the dataset is loaded from. By default, the `WOD_1_1_0_TRAINING` loads up to 128 objects (e.g. vehicles, pedestrians) per scenario. Here, we can save memory and compute by loading only the first 32 objects stored in the scenario.\n",
        "\n",
        "We use the `dataloader.simulator_state_generator` function to create an iterator\n",
        "through Open Motion Dataset scenarios. Calling next on the iterator will retrieve the first scenario in the dataset.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "5-sDgxz9Th1r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "auth.authenticate_user()\n",
        "!gsutil cp gs://waymo_open_dataset_motion_v_1_2_0/uncompressed/tf_example/training/training_tfexample.tfrecord-00000-of-01000 /content/training_tfexample.tfrecord\n"
      ],
      "metadata": {
        "id": "7-XvvMJlQMC7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dkJwTuSLr0gh"
      },
      "outputs": [],
      "source": [
        "\n",
        "config = _config.DatasetConfig(path ='/content/training_tfexample.tfrecord',\n",
        "    data_format=_config.DataFormat.TFRECORD,\n",
        "    max_num_objects=32)\n",
        "data_iter = dataloader.simulator_state_generator(config=config)\n",
        "# Check if the iterator is empty before calling next\n",
        "try:\n",
        "    scenario = next(data_iter)\n",
        "    print(scenario)\n",
        "except StopIteration:\n",
        "    print(\"The data iterator is empty.\")\n",
        "    # Handle empty iterator (e.g., reload data, check config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q1xyeYpLR8J6"
      },
      "source": [
        "Next, we can plot the initial state of this scenario. We use a matplotlib-based visualization available in the `waymax.visualization` package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OY3-OOArsFcU"
      },
      "outputs": [],
      "source": [
        "# Using logged trajectory\n",
        "img = visualization.plot_simulator_state(scenario, use_log_traj=True)\n",
        "mediapy.show_image(img)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H0Z15epRSC23"
      },
      "source": [
        "The Waymo Open Motion Dataset consists of 9-second trajectory snippets. We can visualize the entire logged trajectory as a video as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "06SjvXdRrV3N"
      },
      "outputs": [],
      "source": [
        "imgs = []\n",
        "\n",
        "state = scenario\n",
        "for _ in range(scenario.remaining_timesteps):\n",
        "  state = datatypes.update_state_by_log(state, num_steps=1)\n",
        "  imgs.append(visualization.plot_simulator_state(state, use_log_traj=True))\n",
        "\n",
        "mediapy.show_video(imgs, fps=10)"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "eOpM1SL5SpMH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DINmUYg7y-jI"
      },
      "source": [
        "## Initializing and Running the Simulator\n",
        "\n",
        "Waymax uses a Gym-like interface for running closed-loop simulation.\n",
        "\n",
        "The `env.MultiAgentEnvironment` class defines a stateless simulation interface with the two key methods:\n",
        "- The `reset` method initializes and returns the first simulation state.\n",
        "- The `step` method transitions the simulation and takes as arguments a state and an action and outputs the next state.\n",
        "\n",
        "Crucially, the `MultiAgentEnvironment` does not hold any simulation state itself, and the `reset` and `step` functions have no side effects. This allows us to use functional transforms from JAX, such as using jit compilation to optimize the compuation. It also allows the user to arbitrarily branch and restart simulation from any state, or save the simulation by simply serializing and saving the state object.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Config the multi-agent environment:\n",
        "init_steps = 11\n",
        "\n",
        "# Set the dynamics model the environment is using.\n",
        "# Note each actor interacting with the environment needs to provide action\n",
        "# compatible with this dynamics model.\n",
        "dynamics_model = dynamics.StateDynamics()\n",
        "\n",
        "# Expect users to control all valid object in the scene.\n",
        "env = _env.MultiAgentEnvironment(\n",
        "    dynamics_model=dynamics_model,\n",
        "    config=dataclasses.replace(\n",
        "        _config.EnvironmentConfig(),\n",
        "        max_num_objects=32,\n",
        "        controlled_object=_config.ObjectType.VALID,\n",
        "    ),\n",
        ")"
      ],
      "metadata": {
        "id": "HyvNmNE6T3fH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "876iboHbYx2H"
      },
      "source": [
        "We now create a set of sim agents to run in simulation. By default, the behavior of an object that is not controlled is to replay the behavior stored in the dataset (log playback).\n",
        "\n",
        "For each sim agent, we define the algorithm (such as IDM), and specify which objects the agent controls via the `is_controlled_func`, which is required to return a boolean mask marking which objects are being controlled.\n",
        "\n",
        "The IDM agent we use in this example is the `IDMRoutePolicy`, which follows the spatial trajectory stored in the logs, but adjusts the speed profile based on the IDM rule, which will stop or speed up according to the distance between the vehicle and any objects in front of the vehicle. For the remaining agents, we set them to use a constant speed policy which will follow the logged route with a fixed, constant speed."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Setup a few actors, see visualization below for how each actor behaves.\n",
        "\n",
        "# An actor that doesn't move, controlling all objects with index > 4\n",
        "obj_idx = jnp.arange(32)\n",
        "static_actor = agents.create_constant_speed_actor(\n",
        "    speed=0.0,\n",
        "    dynamics_model=dynamics_model,\n",
        "    is_controlled_func=lambda state: obj_idx > 4,\n",
        ")\n",
        "\n",
        "# IDM actor/policy controlling both object 0 and 1.\n",
        "# Note IDM policy is an actor hard-coded to use dynamics.StateDynamics().\n",
        "actor_0 = agents.IDMRoutePolicy(\n",
        "    is_controlled_func=lambda state: (obj_idx == 0) | (obj_idx == 1)\n",
        ")\n",
        "\n",
        "# Constant speed actor with predefined fixed speed controlling object 2.\n",
        "actor_1 = agents.create_constant_speed_actor(\n",
        "    speed=5.0,\n",
        "    dynamics_model=dynamics_model,\n",
        "    is_controlled_func=lambda state: obj_idx == 2,\n",
        ")\n",
        "\n",
        "# Exper/log actor controlling objects 3 and 4.\n",
        "actor_2 = agents.create_expert_actor(\n",
        "    dynamics_model=dynamics_model,\n",
        "    is_controlled_func=lambda state: (obj_idx == 3) | (obj_idx == 4),\n",
        ")\n",
        "\n",
        "actors = [static_actor, actor_0, actor_1, actor_2]"
      ],
      "metadata": {
        "id": "oY0oHGggUK8O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can (optionally) jit the step and select action functions to speed up computation."
      ],
      "metadata": {
        "id": "cnTTs-H8UT2Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "jit_step = jax.jit(env.step)\n",
        "jit_select_action_list = [jax.jit(actor.select_action) for actor in actors]"
      ],
      "metadata": {
        "id": "iD7G-sSeUWAs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can now write a for loop to all of these agents in simulation together."
      ],
      "metadata": {
        "id": "_YVNshEqUaC0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "states = [env.reset(scenario)]\n",
        "for _ in range(states[0].remaining_timesteps):\n",
        "  current_state = states[-1]\n",
        "\n",
        "  outputs = [\n",
        "      jit_select_action({}, current_state, None, None)\n",
        "      for jit_select_action in jit_select_action_list\n",
        "  ]\n",
        "  action = agents.merge_actions(outputs)\n",
        "  next_state = jit_step(current_state, action)\n",
        "\n",
        "  states.append(next_state)"
      ],
      "metadata": {
        "id": "X_vVS3hxUbdo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s6TmqhLRGc_7"
      },
      "source": [
        "## Visualization of simulation.\n",
        "\n",
        "We can now visualize the result of the simulation loop.\n",
        "\n",
        "On the left side:\n",
        "- Objects 5, 6, and 7 (controlled by static_actor) remain static.\n",
        "- Objects 3 and 4 controlled by log playback, and collide with objects 5 and 6.\n",
        "\n",
        "On the right side:\n",
        "- Object 2 controlled by actor_1 is moving at constant speed 5m/s (i.e. slower than log in this case).\n",
        "- Object 0 and 1, controlled by the IDM agent, follow the log in the beginning, but object 1 slows down when approaching object 2."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "imgs = []\n",
        "for state in states:\n",
        "  imgs.append(visualization.plot_simulator_state(state, use_log_traj=False))\n",
        "mediapy.show_video(imgs, fps=10)"
      ],
      "metadata": {
        "id": "_eaI3hO4UimH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Wayformer Model Setup\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "from typing import Optional, Tuple, List\n",
        "from dataclasses import dataclass"
      ],
      "metadata": {
        "id": "N6qM2ElvUkDx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "g8a-oFncNo86"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@dataclass\n",
        "class WayformerTrainingConfig:\n",
        "    \"\"\"Configuration for Wayformer model training.\"\"\"\n",
        "    num_map_feature: int = 11  # Road feature dimensions\n",
        "    num_agent_feature: int = 9  # Agent feature dimensions\n",
        "    hidden_size: int = 256\n",
        "    max_num_agents: int = 32\n",
        "    num_modes: int = 6\n",
        "    future_len: int = 80  # 8 seconds with 10Hz\n",
        "    past_len: int = 11   # 1 second with 10Hz\n",
        "    dropout: float = 0.1\n",
        "    tx_num_heads: int = 8\n",
        "    max_points_per_lane: int = 40\n",
        "    max_num_roads: int = 50\n",
        "    num_queries_enc: int = 128\n",
        "    num_queries_dec: int = 64\n",
        "    learning_rate: float = 1e-4\n",
        "    batch_size: int = 32\n",
        "    num_epochs: int = 10\n",
        "\n",
        "@dataclass\n",
        "class ModuleOutput:\n",
        "    last_hidden_state: tf.Tensor\n",
        "    kv_cache: Optional[Tuple[tf.Tensor, tf.Tensor]] = None\n",
        "\n",
        "class TrainableQueryProvider(tf.keras.layers.Layer):\n",
        "    def __init__(self, num_queries: int, num_query_channels: int, init_scale: float = 0.02):\n",
        "        super().__init__()\n",
        "        self.num_queries = num_queries\n",
        "        self.num_query_channels = num_query_channels\n",
        "        self.init_scale = init_scale\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        self.query = self.add_weight(\n",
        "            shape=(self.num_queries, self.num_query_channels),\n",
        "            initializer=tf.keras.initializers.RandomNormal(stddev=self.init_scale),\n",
        "            trainable=True,\n",
        "            name='query'\n",
        "        )\n",
        "\n",
        "    def call(self, x=None):\n",
        "        return tf.expand_dims(self.query, 0)  # Add batch dimension\n",
        "\n",
        "class MultiHeadAttention(tf.keras.layers.Layer):\n",
        "    def __init__(\n",
        "        self,\n",
        "        num_heads: int,\n",
        "        num_q_input_channels: int,\n",
        "        num_kv_input_channels: int,\n",
        "        num_qk_channels: Optional[int] = None,\n",
        "        num_v_channels: Optional[int] = None,\n",
        "        num_output_channels: Optional[int] = None,\n",
        "        max_heads_parallel: Optional[int] = None,\n",
        "        causal_attention: bool = False,\n",
        "        dropout: float = 0.0,\n",
        "        qkv_bias: bool = True,\n",
        "        out_bias: bool = True,\n",
        "    ):\n",
        "        super().__init__()\n",
        "\n",
        "        if num_qk_channels is None:\n",
        "            num_qk_channels = num_q_input_channels\n",
        "\n",
        "        if num_v_channels is None:\n",
        "            num_v_channels = num_qk_channels\n",
        "\n",
        "        if num_output_channels is None:\n",
        "            num_output_channels = num_q_input_channels\n",
        "\n",
        "        self.num_heads = num_heads\n",
        "        self.dp_scale = (num_qk_channels // num_heads) ** -0.5\n",
        "        self.causal_attention = causal_attention\n",
        "\n",
        "        self.q_proj = tf.keras.layers.Dense(num_qk_channels, use_bias=qkv_bias)\n",
        "        self.k_proj = tf.keras.layers.Dense(num_qk_channels, use_bias=qkv_bias)\n",
        "        self.v_proj = tf.keras.layers.Dense(num_v_channels, use_bias=qkv_bias)\n",
        "        self.o_proj = tf.keras.layers.Dense(num_output_channels, use_bias=out_bias)\n",
        "        self.dropout = tf.keras.layers.Dropout(dropout)\n",
        "\n",
        "    def call(self, x_q, x_kv, pad_mask=None, training=False):\n",
        "        batch_size = tf.shape(x_q)[0]\n",
        "\n",
        "        # Linear projections and reshape for multi-head attention\n",
        "        q = self.q_proj(x_q)  # [batch_size, seq_len_q, d_model]\n",
        "        k = self.k_proj(x_kv)  # [batch_size, seq_len_k, d_model]\n",
        "        v = self.v_proj(x_kv)  # [batch_size, seq_len_k, d_model]\n",
        "\n",
        "        # Reshape to [batch_size, num_heads, seq_len, depth]\n",
        "        q = self._reshape_for_heads(q)\n",
        "        k = self._reshape_for_heads(k)\n",
        "        v = self._reshape_for_heads(v)\n",
        "\n",
        "        # Scale query\n",
        "        q = q * self.dp_scale\n",
        "\n",
        "        # Calculate attention scores\n",
        "        attn = tf.matmul(q, k, transpose_b=True)\n",
        "\n",
        "        if pad_mask is not None:\n",
        "            pad_mask = tf.expand_dims(tf.expand_dims(pad_mask, 1), 1)\n",
        "            attn = tf.where(pad_mask, tf.float32.min, attn)\n",
        "\n",
        "        if self.causal_attention:\n",
        "            causal_mask = self._create_causal_mask(tf.shape(q)[2], tf.shape(k)[2])\n",
        "            attn = tf.where(causal_mask, tf.float32.min, attn)\n",
        "\n",
        "        attn = tf.nn.softmax(attn, axis=-1)\n",
        "        attn = self.dropout(attn, training=training)\n",
        "\n",
        "        # Calculate output\n",
        "        output = tf.matmul(attn, v)\n",
        "        output = self._reshape_from_heads(output)\n",
        "        output = self.o_proj(output)\n",
        "\n",
        "        return ModuleOutput(last_hidden_state=output)\n",
        "\n",
        "    def _reshape_for_heads(self, x):\n",
        "        batch_size = tf.shape(x)[0]\n",
        "        seq_len = tf.shape(x)[1]\n",
        "        depth = tf.shape(x)[2] // self.num_heads\n",
        "\n",
        "        x = tf.reshape(x, [batch_size, seq_len, self.num_heads, depth])\n",
        "        return tf.transpose(x, [0, 2, 1, 3])\n",
        "\n",
        "    def _reshape_from_heads(self, x):\n",
        "        batch_size = tf.shape(x)[0]\n",
        "        seq_len = tf.shape(x)[2]\n",
        "\n",
        "        x = tf.transpose(x, [0, 2, 1, 3])\n",
        "        return tf.reshape(x, [batch_size, seq_len, -1])\n",
        "\n",
        "    def _create_causal_mask(self, seq_len_q, seq_len_k):\n",
        "        mask = tf.ones((seq_len_q, seq_len_k), dtype=tf.bool)\n",
        "        mask = tf.linalg.band_part(mask, -1, 0)  # Lower triangular\n",
        "        return mask\n",
        "\n",
        "class PerceiverEncoder(tf.keras.layers.Layer):\n",
        "    def __init__(\n",
        "        self,\n",
        "        num_latents: int,\n",
        "        num_latent_channels: int,\n",
        "        num_cross_attention_heads: int = 4,\n",
        "        num_cross_attention_qk_channels: Optional[int] = None,\n",
        "        num_cross_attention_v_channels: Optional[int] = None,\n",
        "        num_cross_attention_layers: int = 1,\n",
        "        dropout: float = 0.1,\n",
        "        init_scale: float = 0.02,\n",
        "    ):\n",
        "        super().__init__()\n",
        "\n",
        "        self.latent_provider = TrainableQueryProvider(\n",
        "            num_latents,\n",
        "            num_latent_channels,\n",
        "            init_scale=init_scale\n",
        "        )\n",
        "\n",
        "        self.cross_attention = MultiHeadAttention(\n",
        "            num_heads=num_cross_attention_heads,\n",
        "            num_q_input_channels=num_latent_channels,\n",
        "            num_kv_input_channels=num_latent_channels,\n",
        "            num_qk_channels=num_cross_attention_qk_channels,\n",
        "            num_v_channels=num_cross_attention_v_channels,\n",
        "            dropout=dropout\n",
        "        )\n",
        "\n",
        "        self.self_attention = MultiHeadAttention(\n",
        "            num_heads=num_cross_attention_heads,\n",
        "            num_q_input_channels=num_latent_channels,\n",
        "            num_kv_input_channels=num_latent_channels,\n",
        "            dropout=dropout\n",
        "        )\n",
        "\n",
        "        self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.dropout = tf.keras.layers.Dropout(dropout)\n",
        "\n",
        "    def call(self, x, pad_mask=None, training=False):\n",
        "        x_latent = self.latent_provider()\n",
        "\n",
        "        # Cross attention\n",
        "        residual = x_latent\n",
        "        x_latent = self.layer_norm1(x_latent)\n",
        "        cross_attn_output = self.cross_attention(x_latent, x, pad_mask=pad_mask, training=training)\n",
        "        x_latent = residual + self.dropout(cross_attn_output.last_hidden_state, training=training)\n",
        "\n",
        "        # Self attention\n",
        "        residual = x_latent\n",
        "        x_latent = self.layer_norm2(x_latent)\n",
        "        self_attn_output = self.self_attention(x_latent, x_latent, training=training)\n",
        "        x_latent = residual + self.dropout(self_attn_output.last_hidden_state, training=training)\n",
        "\n",
        "        return x_latent\n",
        "\n",
        "\n",
        "def train_wayformer(model_config, action_space, training_data):\n",
        "    \"\"\"Train Wayformer model with better error handling and debugging\"\"\"\n",
        "    model = Wayformer(model_config, action_space)\n",
        "    optimizer = tf.keras.optimizers.Adam(learning_rate=model_config.learning_rate)\n",
        "\n",
        "    # Convert data to numpy arrays for TensorFlow\n",
        "    training_data = [\n",
        "        {k: (v.numpy() if hasattr(v, 'numpy') else v) for k, v in scenario.items()}\n",
        "        for scenario in training_data\n",
        "    ]\n",
        "\n",
        "    num_samples = len(training_data)\n",
        "    num_batches = num_samples // model_config.batch_size\n",
        "    print(f\"Training on {num_samples} samples with {num_batches} batches per epoch\")\n",
        "\n",
        "    for epoch in range(model_config.num_epochs):\n",
        "        print(f\"\\nEpoch {epoch + 1}/{model_config.num_epochs}\")\n",
        "        total_loss = 0.0\n",
        "        batch_losses = []\n",
        "\n",
        "        # Shuffle the data\n",
        "        np.random.shuffle(training_data)\n",
        "\n",
        "        for batch_idx in tqdm(range(num_batches)):\n",
        "            try:\n",
        "                # Get batch\n",
        "                start_idx = batch_idx * model_config.batch_size\n",
        "                end_idx = start_idx + model_config.batch_size\n",
        "                batch_data = training_data[start_idx:end_idx]\n",
        "\n",
        "                # Prepare input\n",
        "                model_inputs = prepare_model_input(batch_data)\n",
        "                model_inputs = {k: tf.convert_to_tensor(v, dtype=tf.float32) for k, v in model_inputs.items()}\n",
        "\n",
        "                # Debug shapes\n",
        "                print(f\"\\nBatch {batch_idx} input shapes:\")\n",
        "                for k, v in model_inputs.items():\n",
        "                    print(f\"{k}: {v.shape}\")\n",
        "\n",
        "                with tf.GradientTape() as tape:\n",
        "                    predictions = model(model_inputs, training=True)\n",
        "                    loss = compute_trajectory_loss(\n",
        "                        predictions['predicted_trajectory'],\n",
        "                        predictions['predicted_probability'],\n",
        "                        model_inputs['future_trajectory'],\n",
        "                        tf.ones_like(model_inputs['future_trajectory'][..., 0])\n",
        "                    )\n",
        "\n",
        "                # Compute and apply gradients\n",
        "                gradients = tape.gradient(loss, model.trainable_variables)\n",
        "                optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "\n",
        "                batch_losses.append(float(loss.numpy()))\n",
        "                total_loss += float(loss.numpy())\n",
        "\n",
        "                print(f\"Batch {batch_idx} loss: {float(loss.numpy()):.4f}\")\n",
        "\n",
        "            except Exception as e:\n",
        "                print(f\"\\nError processing batch {batch_idx}:\")\n",
        "                print(str(e))\n",
        "                print(\"\\nFull traceback:\")\n",
        "                import traceback\n",
        "                traceback.print_exc()\n",
        "                continue\n",
        "\n",
        "        # Print epoch statistics\n",
        "        if batch_losses:\n",
        "            avg_loss = total_loss / len(batch_losses)\n",
        "            print(f\"\\nEpoch {epoch + 1} statistics:\")\n",
        "            print(f\"Average Loss: {avg_loss:.4f}\")\n",
        "            print(f\"Min Batch Loss: {min(batch_losses):.4f}\")\n",
        "            print(f\"Max Batch Loss: {max(batch_losses):.4f}\")\n",
        "        else:\n",
        "            print(f\"\\nEpoch {epoch + 1}: No successful batches\")\n",
        "\n",
        "    return model"
      ],
      "metadata": {
        "id": "YX5nfz3CHxL1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def process_waymax_data(state, action_space):\n",
        "    \"\"\"Convert waymax SimulatorState to model input format with consistent shapes\"\"\"\n",
        "    trajectory = state.sim_trajectory\n",
        "    timestep = state.timestep\n",
        "    past_steps = 11  # Match the config.past_len\n",
        "    future_steps = 80  # Match the config.future_len\n",
        "\n",
        "    # Get past trajectory features (last past_steps timesteps before current)\n",
        "    past_idx = slice(max(0, timestep - past_steps + 1), timestep + 1)\n",
        "\n",
        "    # Extract and pad past features if necessary\n",
        "    xy = trajectory.xy[..., past_idx, :]\n",
        "    vel_xy = trajectory.vel_xy[..., past_idx, :]\n",
        "    yaw = trajectory.yaw[..., past_idx]\n",
        "    length = trajectory.length[..., past_idx]\n",
        "    width = trajectory.width[..., past_idx]\n",
        "    height = trajectory.height[..., past_idx]\n",
        "    valid = trajectory.valid[..., past_idx]\n",
        "\n",
        "    # Pad if we don't have enough past steps\n",
        "    if xy.shape[-2] < past_steps:\n",
        "        pad_length = past_steps - xy.shape[-2]\n",
        "        xy = jnp.pad(xy, ((0, 0), (pad_length, 0), (0, 0)), mode='edge')\n",
        "        vel_xy = jnp.pad(vel_xy, ((0, 0), (pad_length, 0), (0, 0)), mode='edge')\n",
        "        yaw = jnp.pad(yaw, ((0, 0), (pad_length, 0)), mode='edge')\n",
        "        length = jnp.pad(length, ((0, 0), (pad_length, 0)), mode='edge')\n",
        "        width = jnp.pad(width, ((0, 0), (pad_length, 0)), mode='edge')\n",
        "        height = jnp.pad(height, ((0, 0), (pad_length, 0)), mode='edge')\n",
        "        valid = jnp.pad(valid, ((0, 0), (pad_length, 0)), mode='constant', constant_values=0)\n",
        "\n",
        "    # Expand dimensions for single-value features\n",
        "    yaw = jnp.expand_dims(yaw, axis=-1)\n",
        "    length = jnp.expand_dims(length, axis=-1)\n",
        "    width = jnp.expand_dims(width, axis=-1)\n",
        "    height = jnp.expand_dims(height, axis=-1)\n",
        "    valid = jnp.expand_dims(valid, axis=-1)\n",
        "\n",
        "    # Broadcast is_sdc to match time dimension\n",
        "    is_sdc = jnp.expand_dims(state.object_metadata.is_sdc, axis=(1, -1))\n",
        "    is_sdc = jnp.repeat(is_sdc, past_steps, axis=1)\n",
        "\n",
        "    # Concatenate all features\n",
        "    current_features = jnp.concatenate([\n",
        "        xy,                     # (..., past_steps, 2)\n",
        "        vel_xy,                 # (..., past_steps, 2)\n",
        "        yaw,                    # (..., past_steps, 1)\n",
        "        length,                 # (..., past_steps, 1)\n",
        "        width,                  # (..., past_steps, 1)\n",
        "        height,                 # (..., past_steps, 1)\n",
        "        is_sdc,                 # (..., past_steps, 1)\n",
        "        valid                   # (..., past_steps, 1)\n",
        "    ], axis=-1)\n",
        "\n",
        "    # Get future trajectory features\n",
        "    future_idx = slice(timestep + 1, min(timestep + 1 + future_steps, trajectory.num_timesteps))\n",
        "    future_xy = trajectory.xy[..., future_idx, :]\n",
        "    future_vel_xy = trajectory.vel_xy[..., future_idx, :]\n",
        "    future_yaw = jnp.expand_dims(trajectory.yaw[..., future_idx], axis=-1)\n",
        "    future_valid = jnp.expand_dims(trajectory.valid[..., future_idx], axis=-1)\n",
        "\n",
        "    # Pad future features if necessary\n",
        "    if future_xy.shape[-2] < future_steps:\n",
        "        pad_length = future_steps - future_xy.shape[-2]\n",
        "        future_xy = jnp.pad(future_xy, ((0, 0), (0, pad_length), (0, 0)), mode='edge')\n",
        "        future_vel_xy = jnp.pad(future_vel_xy, ((0, 0), (0, pad_length), (0, 0)), mode='edge')\n",
        "        future_yaw = jnp.pad(future_yaw, ((0, 0), (0, pad_length), (0, 0)), mode='edge')\n",
        "        future_valid = jnp.pad(future_valid, ((0, 0), (0, pad_length), (0, 0)), mode='constant', constant_values=0)\n",
        "\n",
        "    future_features = jnp.concatenate([\n",
        "        future_xy,\n",
        "        future_vel_xy,\n",
        "        future_yaw,\n",
        "        future_valid\n",
        "    ], axis=-1)\n",
        "\n",
        "    # Process road features\n",
        "    if hasattr(state, 'roadgraph_points') and state.roadgraph_points is not None:\n",
        "        road_xyz = state.roadgraph_points.xyz\n",
        "        road_dir = jnp.stack([\n",
        "            state.roadgraph_points.dir_x,\n",
        "            state.roadgraph_points.dir_y,\n",
        "            state.roadgraph_points.dir_z\n",
        "        ], axis=-1)\n",
        "        road_valid = jnp.expand_dims(state.roadgraph_points.valid, axis=-1)\n",
        "        road_features = jnp.concatenate([road_xyz, road_dir, road_valid], axis=-1)\n",
        "\n",
        "        # Truncate or pad to fixed size\n",
        "        max_road_points = 1024\n",
        "        if road_features.shape[0] > max_road_points:\n",
        "            road_features = road_features[:max_road_points]\n",
        "        else:\n",
        "            pad_length = max_road_points - road_features.shape[0]\n",
        "            road_features = jnp.pad(road_features, ((0, pad_length), (0, 0)), mode='constant')\n",
        "    else:\n",
        "        road_features = jnp.zeros((1024, 7))\n",
        "    # Assume ego vehicle is at index 0\n",
        "    ego_index = 0  # Update this if necessary based on your data\n",
        "    # Extract ego's past and future data\n",
        "    ego_past_positions = xy[ego_index]  # Shape: (past_steps, 2)\n",
        "    ego_past_velocities = vel_xy[ego_index]  # Shape: (past_steps, 2)\n",
        "    ego_future_positions = future_xy[ego_index]  # Shape: (future_steps, 2)\n",
        "    ego_future_velocities = future_vel_xy[ego_index]  # Shape: (future_steps, 2)\n",
        "    # Combine past and future positions and velocities\n",
        "    positions = np.concatenate([ego_past_positions, ego_future_positions], axis=0)\n",
        "    velocities = np.concatenate([ego_past_velocities, ego_future_velocities], axis=0)\n",
        "\n",
        "    # Encode the future trajectory into motion tokens\n",
        "    # Note: For encoding, we need positions and velocities at T+1 points to compute T accelerations\n",
        "    token_indices = encode_trajectory(positions, velocities, action_space)\n",
        "\n",
        "    return {\n",
        "        'current_features': current_features,  # (32, past_steps=11, 10)\n",
        "        'future_features': future_features,    # (32, future_steps=80, 6)\n",
        "        'road_features': road_features,        # (1024, 7)\n",
        "        'object_metadata': {\n",
        "            'is_sdc': state.object_metadata.is_sdc,\n",
        "            'object_types': state.object_metadata.object_types,\n",
        "            'is_valid': state.object_metadata.is_valid,\n",
        "        },\n",
        "        'motion_tokens': token_indices[past_steps - 1:]  # Use future tokens only\n",
        "    }\n",
        "\n",
        "def prepare_model_input(batch_data):\n",
        "    \"\"\"Prepare batch data for model input with correct shapes\"\"\"\n",
        "    batch_size = len(batch_data)\n",
        "\n",
        "    # Stack all features\n",
        "    batched_data = {\n",
        "        'current_features': jnp.stack([d['current_features'] for d in batch_data]),\n",
        "        'future_features': jnp.stack([d['future_features'] for d in batch_data]),\n",
        "        'road_features': jnp.stack([d['road_features'] for d in batch_data]),\n",
        "        'object_metadata': {\n",
        "            'is_sdc': jnp.stack([d['object_metadata']['is_sdc'] for d in batch_data])\n",
        "        },\n",
        "        'motion_tokens': jnp.stack([d['motion_tokens'] for d in batch_data])\n",
        "    }\n",
        "\n",
        "    # Find ego vehicle (SDC) in each scenario\n",
        "    ego_indices = jnp.argmax(batched_data['object_metadata']['is_sdc'], axis=-1)\n",
        "\n",
        "    # Extract ego and other agent features\n",
        "    ego_features = []\n",
        "    ego_future_features = []\n",
        "    other_agents = []\n",
        "\n",
        "    for i in range(batch_size):\n",
        "        # Get ego features\n",
        "        ego_feat = batched_data['current_features'][i, ego_indices[i]]\n",
        "        ego_features.append(ego_feat)\n",
        "\n",
        "        # Get ego future features\n",
        "        ego_future_feat = batched_data['future_features'][i, ego_indices[i]]\n",
        "        ego_future_features.append(ego_future_feat)\n",
        "\n",
        "        # Get other agent features (excluding ego)\n",
        "        agents = []\n",
        "        for j in range(batched_data['current_features'].shape[1]):\n",
        "            if j != ego_indices[i]:\n",
        "                agents.append(batched_data['current_features'][i, j])\n",
        "        agents = jnp.stack(agents[:31])  # Limit to 31 other agents\n",
        "\n",
        "        # Pad if we have fewer than 31 agents\n",
        "        if agents.shape[0] < 31:\n",
        "            pad_length = 31 - agents.shape[0]\n",
        "            agents = jnp.pad(agents, ((0, pad_length), (0, 0), (0, 0)), mode='constant')\n",
        "\n",
        "        other_agents.append(agents)\n",
        "\n",
        "    return {\n",
        "        'ego_in': jnp.stack(ego_features),          # (batch_size, past_steps, features)\n",
        "        'agents_in': jnp.stack(other_agents),       # (batch_size, 31, past_steps, features)\n",
        "        'roads': batched_data['road_features'],     # (batch_size, max_road_points, features)\n",
        "        'future_trajectory': jnp.stack(ego_future_features),  # (batch_size, future_steps, features)\n",
        "        'motion_tokens': batched_data['motion_tokens']  # (batch_size, future_steps)\n",
        "    }\n",
        "\n",
        "def create_training_dataset(config, action_space, num_scenarios=1000):\n",
        "    \"\"\"Create training dataset from waymax data\"\"\"\n",
        "    data_iter = dataloader.simulator_state_generator(config=config)\n",
        "    dataset = []\n",
        "\n",
        "    for i in tqdm(range(num_scenarios)):\n",
        "        try:\n",
        "            scenario = next(data_iter)\n",
        "            if jnp.any(scenario.object_metadata.is_valid):\n",
        "                try:\n",
        "                    processed_data = process_waymax_data(scenario,action_space)\n",
        "                    dataset.append(processed_data)\n",
        "                    if i == 0:  # Print shapes for first successful scenario\n",
        "                        print(\"\\nFirst successful scenario shapes:\")\n",
        "                        print(f\"current_features: {processed_data['current_features'].shape}\")\n",
        "                        print(f\"future_features: {processed_data['future_features'].shape}\")\n",
        "                        print(f\"road_features: {processed_data['road_features'].shape}\")\n",
        "                except Exception as e:\n",
        "                    print(f\"Failed to process scenario: {str(e)}\")\n",
        "                    continue\n",
        "        except StopIteration:\n",
        "            break\n",
        "\n",
        "    print(f\"\\nSuccessfully processed {len(dataset)} scenarios\")\n",
        "    return dataset"
      ],
      "metadata": {
        "id": "xyvX9GGmI_FV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def compute_trajectory_loss(pred_trajectories, pred_probabilities, target_trajectories, is_valid):\n",
        "    \"\"\"Compute training loss for trajectory prediction with proper shape handling\"\"\"\n",
        "    # pred_trajectories shape: [batch, num_modes, time, 5]\n",
        "    # target_trajectories shape: [batch, time, 6]\n",
        "    # is_valid shape: [batch, time]\n",
        "\n",
        "    batch_size = tf.shape(pred_trajectories)[0]\n",
        "    num_modes = tf.shape(pred_trajectories)[1]\n",
        "\n",
        "    # Expand target trajectories for comparison with each mode\n",
        "    target_expanded = tf.expand_dims(target_trajectories[..., :5], axis=1)  # [batch, 1, time, 5]\n",
        "    target_expanded = tf.tile(target_expanded, [1, num_modes, 1, 1])  # [batch, num_modes, time, 5]\n",
        "\n",
        "    # Compute displacement error for each mode\n",
        "    displacement_error = tf.reduce_mean(\n",
        "        tf.sqrt(tf.reduce_sum(\n",
        "            tf.square(pred_trajectories - target_expanded),\n",
        "            axis=-1\n",
        "        )),\n",
        "        axis=-1  # Average over time\n",
        "    )  # [batch, num_modes]\n",
        "\n",
        "    # Find best matching mode for each sample\n",
        "    min_mode_error = tf.reduce_min(displacement_error, axis=1)  # [batch]\n",
        "\n",
        "    # Compute probability loss\n",
        "    prob_loss = -tf.math.log(tf.clip_by_value(pred_probabilities, 1e-7, 1.0))\n",
        "    prob_loss = tf.reduce_mean(prob_loss)  # Scalar\n",
        "\n",
        "    # Combine losses\n",
        "    total_loss = tf.reduce_mean(min_mode_error) + 0.1 * prob_loss\n",
        "\n",
        "    return total_loss\n",
        "\n",
        "class TransformerDecoderLayer(tf.keras.layers.Layer):\n",
        "    def __init__(self, d_model, num_heads, dff, dropout_rate=0.1):\n",
        "        super(TransformerDecoderLayer, self).__init__()\n",
        "\n",
        "        # Self-attention (masked)\n",
        "        self.mha1 = tf.keras.layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=d_model // num_heads, dropout=dropout_rate)\n",
        "\n",
        "        # Cross-attention with encoder outputs\n",
        "        self.mha2 = tf.keras.layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=d_model // num_heads, dropout=dropout_rate)\n",
        "\n",
        "        # Point-wise feedforward network\n",
        "        self.ffn = tf.keras.Sequential([\n",
        "            tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)\n",
        "            tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
        "        ])\n",
        "\n",
        "        # Layer normalizations\n",
        "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "\n",
        "        # Dropouts\n",
        "        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)\n",
        "        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)\n",
        "        self.dropout3 = tf.keras.layers.Dropout(dropout_rate)\n",
        "\n",
        "    def call(self, x, enc_output, look_ahead_mask=None, padding_mask=None, training=False):\n",
        "        # x.shape == (batch_size, target_seq_len, d_model)\n",
        "        # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
        "\n",
        "        # Masked self-attention (decoder's self-attention)\n",
        "        attn1 = self.mha1(\n",
        "            query=x,\n",
        "            key=x,\n",
        "            value=x,\n",
        "            attention_mask=look_ahead_mask,\n",
        "            training=training\n",
        "        )  # (batch_size, target_seq_len, d_model)\n",
        "        attn1 = self.dropout1(attn1, training=training)\n",
        "        out1 = self.layernorm1(x + attn1)  # Residual connection and layer norm\n",
        "\n",
        "        # Cross-attention with encoder outputs\n",
        "        attn2 = self.mha2(\n",
        "            query=out1,\n",
        "            key=enc_output,\n",
        "            value=enc_output,\n",
        "            attention_mask=padding_mask,\n",
        "            training=training\n",
        "        )  # (batch_size, target_seq_len, d_model)\n",
        "        attn2 = self.dropout2(attn2, training=training)\n",
        "        out2 = self.layernorm2(out1 + attn2)  # Residual connection and layer norm\n",
        "\n",
        "        # Feedforward network\n",
        "        ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n",
        "        ffn_output = self.dropout3(ffn_output, training=training)\n",
        "        out3 = self.layernorm3(out2 + ffn_output)  # Residual connection and layer norm\n",
        "\n",
        "        return out3\n",
        "\n",
        "class Wayformer(tf.keras.Model):\n",
        "    def __init__(self, config, action_space):\n",
        "        super().__init__()\n",
        "        # Initialize dimensions and parameters\n",
        "        self.map_attr = config.num_map_feature\n",
        "        self.k_attr = config.num_agent_feature\n",
        "        self.d_k = config.hidden_size\n",
        "        self._M = config.max_num_agents\n",
        "        self.c = config.num_modes\n",
        "        self.T = config.future_len\n",
        "        self.dropout = config.dropout\n",
        "        self.num_heads = config.tx_num_heads\n",
        "        self.past_T = config.past_len\n",
        "\n",
        "        # Input encoders\n",
        "        self.road_pts_lin = tf.keras.layers.Dense(self.d_k)\n",
        "        self.agents_dynamic_encoder = tf.keras.layers.Dense(self.d_k)\n",
        "\n",
        "        # Positional embeddings\n",
        "        self.temporal_embedding = self.add_weight(\n",
        "            shape=(self.past_T, self.d_k),\n",
        "            initializer='zeros',\n",
        "            trainable=True,\n",
        "            name='temporal_emb'\n",
        "        )\n",
        "\n",
        "        self.agent_embedding = self.add_weight(\n",
        "            shape=(32, self.d_k),  # 32 = 1 ego + 31 other agents\n",
        "            initializer='zeros',\n",
        "            trainable=True,\n",
        "            name='agent_emb'\n",
        "        )\n",
        "\n",
        "        # Agent processing layers\n",
        "        self.agent_encoder = tf.keras.Sequential([\n",
        "            tf.keras.layers.Dense(self.d_k, activation='relu'),\n",
        "            tf.keras.layers.LayerNormalization(),\n",
        "            tf.keras.layers.Dropout(self.dropout)\n",
        "        ])\n",
        "\n",
        "        self.agent_transformer = tf.keras.layers.MultiHeadAttention(\n",
        "            num_heads=8,\n",
        "            key_dim=self.d_k//8,\n",
        "            dropout=self.dropout\n",
        "        )\n",
        "\n",
        "        # Road processing layers\n",
        "        self.road_encoder = tf.keras.Sequential([\n",
        "            tf.keras.layers.Dense(self.d_k, activation='relu'),\n",
        "            tf.keras.layers.LayerNormalization(),\n",
        "            tf.keras.layers.Dropout(self.dropout)\n",
        "        ])\n",
        "\n",
        "        # Output processing\n",
        "        self.perceiver = PerceiverEncoder(\n",
        "            num_latents=config.num_queries_enc,\n",
        "            num_latent_channels=self.d_k\n",
        "        )\n",
        "\n",
        "        self.output_query = self.add_weight(\n",
        "            shape=(self.c, self.d_k),\n",
        "            initializer='random_normal',\n",
        "            trainable=True,\n",
        "            name='output_query'\n",
        "        )\n",
        "\n",
        "        self.trajectory_projection = tf.keras.layers.Dense(5 * self.T)\n",
        "        self.mode_projection = tf.keras.layers.Dense(1)\n",
        "\n",
        "        self.action_space = action_space\n",
        "        self.vocab_size = action_space.vocab_size\n",
        "\n",
        "        # Token embedding for motion tokens\n",
        "        self.token_embedding = tf.keras.layers.Embedding(\n",
        "            input_dim=self.vocab_size, output_dim=self.d_k)\n",
        "\n",
        "        # Positional encoding for tokens\n",
        "        self.positional_encoding = self.add_weight(\n",
        "            shape=(config.future_len, self.d_k),\n",
        "            initializer='zeros',\n",
        "            trainable=True,\n",
        "            name='token_positional_encoding'\n",
        "        )\n",
        "\n",
        "        # Transformer decoder layers\n",
        "        self.decoder_layers = [\n",
        "            TransformerDecoderLayer(\n",
        "                d_model=self.d_k,\n",
        "                num_heads=config.tx_num_heads,\n",
        "                dff=1024,\n",
        "                dropout_rate=config.dropout\n",
        "            ) for _ in range(4)\n",
        "        ]\n",
        "\n",
        "        # Output projection to predict token logits\n",
        "        self.token_projection = tf.keras.layers.Dense(self.vocab_size)\n",
        "\n",
        "    def build(self, input_shapes):\n",
        "        \"\"\"Build the model based on input shapes\"\"\"\n",
        "        ego_shape = input_shapes['ego_in']\n",
        "        agents_shape = input_shapes['agents_in']\n",
        "        roads_shape = input_shapes['roads']\n",
        "\n",
        "        self.built = True\n",
        "\n",
        "    def encode_agents(self, ego_in, agents_in, training=False):\n",
        "        \"\"\"Process agent features with shape tracking\"\"\"\n",
        "        batch_size = tf.shape(ego_in)[0]\n",
        "\n",
        "        # Print shapes for debugging\n",
        "        print(f\"ego_in shape: {ego_in.shape}\")\n",
        "        print(f\"agents_in shape: {agents_in.shape}\")\n",
        "\n",
        "        # Reshape ego to add agent dimension\n",
        "        ego_expanded = tf.expand_dims(ego_in, axis=1)  # [batch, 1, time, features]\n",
        "\n",
        "        # Concatenate ego with other agents\n",
        "        all_agents = tf.concat([ego_expanded, agents_in], axis=1)  # [batch, 32, time, features]\n",
        "        print(f\"all_agents shape after concat: {all_agents.shape}\")\n",
        "\n",
        "        # Encode agent features\n",
        "        encoded = self.agent_encoder(all_agents, training=training)\n",
        "        print(f\"encoded shape after agent_encoder: {encoded.shape}\")\n",
        "\n",
        "        # Add temporal embeddings\n",
        "        temporal_emb = tf.expand_dims(tf.expand_dims(self.temporal_embedding, 0), 0)\n",
        "        temporal_emb = tf.tile(temporal_emb, [batch_size, encoded.shape[1], 1, 1])\n",
        "        encoded = encoded + temporal_emb\n",
        "\n",
        "        # Add agent embeddings\n",
        "        agent_emb = tf.expand_dims(tf.expand_dims(self.agent_embedding, 0), 2)\n",
        "        agent_emb = tf.tile(agent_emb, [batch_size, 1, encoded.shape[2], 1])\n",
        "        encoded = encoded + agent_emb\n",
        "\n",
        "        print(f\"encoded shape before transformer: {encoded.shape}\")\n",
        "\n",
        "        # Apply self-attention\n",
        "        attended = self.agent_transformer(\n",
        "            query=encoded,\n",
        "            key=encoded,\n",
        "            value=encoded,\n",
        "            training=training\n",
        "        )\n",
        "        print(f\"attended shape after transformer: {attended.shape}\")\n",
        "\n",
        "        # Final reshape\n",
        "        final_encoded = tf.reshape(attended, [batch_size, -1, self.d_k])\n",
        "        print(f\"final_encoded shape: {final_encoded.shape}\")\n",
        "\n",
        "        return final_encoded\n",
        "\n",
        "    def encode_roads(self, roads, training=False):\n",
        "        # Process road features\n",
        "        encoded = self.road_encoder(roads, training=training)\n",
        "        return tf.reshape(encoded, [tf.shape(encoded)[0], -1, self.d_k])\n",
        "\n",
        "    def call(self, inputs, training=False):\n",
        "        \"\"\"Forward pass with shape debugging\"\"\"\n",
        "        print(\"\\nInput shapes:\")\n",
        "        for k, v in inputs.items():\n",
        "            print(f\"{k}: {v.shape}\")\n",
        "\n",
        "        ego_in = inputs['ego_in']\n",
        "        agents_in = inputs['agents_in']\n",
        "        roads = inputs['roads']\n",
        "\n",
        "        # Process agents\n",
        "        agents_encoded = self.encode_agents(ego_in, agents_in, training)\n",
        "        print(f\"agents_encoded shape: {agents_encoded.shape}\")\n",
        "\n",
        "        # Process roads\n",
        "        roads_encoded = self.encode_roads(roads, training)\n",
        "        print(f\"roads_encoded shape: {roads_encoded.shape}\")\n",
        "\n",
        "        # Combine features\n",
        "        combined_features = tf.concat([agents_encoded, roads_encoded], axis=1)\n",
        "        print(f\"combined_features shape: {combined_features.shape}\")\n",
        "\n",
        "        # Apply perceiver encoding\n",
        "        context = self.perceiver(combined_features, training=training)\n",
        "        print(f\"context shape: {context.shape}\")\n",
        "\n",
        "        # Generate outputs\n",
        "        batch_size = tf.shape(ego_in)[0]\n",
        "        query = tf.tile(tf.expand_dims(self.output_query, 0), [batch_size, 1, 1])\n",
        "        print(f\"query shape: {query.shape}\")\n",
        "\n",
        "        # Get trajectory predictions\n",
        "        output_features = tf.matmul(query, context, transpose_b=True)\n",
        "        print(f\"output_features shape: {output_features.shape}\")\n",
        "\n",
        "        trajectories = self.trajectory_projection(output_features)\n",
        "        trajectories = tf.reshape(trajectories, [batch_size, self.c, self.T, -1])\n",
        "        print(f\"trajectories shape: {trajectories.shape}\")\n",
        "\n",
        "        # Get mode probabilities\n",
        "        mode_logits = tf.squeeze(self.mode_projection(output_features), axis=-1)\n",
        "        mode_probs = tf.nn.softmax(mode_logits, axis=-1)\n",
        "        print(f\"mode_probs shape: {mode_probs.shape}\")\n",
        "\n",
        "        motion_tokens = inputs['motion_tokens']  # Shape: (batch_size, future_steps)\n",
        "\n",
        "        # Embed motion tokens\n",
        "        token_embeddings = self.token_embedding(motion_tokens)  # Shape: (batch_size, future_steps, d_k)\n",
        "        token_embeddings += self.positional_encoding  # Add positional encoding\n",
        "\n",
        "        # Prepare attention masks\n",
        "        look_ahead_mask = tf.linalg.band_part(tf.ones((self.T, self.T)), -1, 0)  # Causal mask\n",
        "\n",
        "        # Pass through decoder layers\n",
        "        decoder_output = token_embeddings\n",
        "        for layer in self.decoder_layers:\n",
        "            decoder_output = layer(\n",
        "                decoder_output, context, training=training,\n",
        "                look_ahead_mask=look_ahead_mask\n",
        "            )\n",
        "\n",
        "        # Predict logits for next tokens\n",
        "        logits = self.token_projection(decoder_output)  # Shape: (batch_size, future_steps, vocab_size)\n",
        "\n",
        "        return {\n",
        "            'predicted_trajectory': trajectories,\n",
        "            'predicted_probability': mode_probs,\n",
        "            'scene_emb': tf.reshape(output_features, [batch_size, -1]),\n",
        "            'logits': logits\n",
        "        }"
      ],
      "metadata": {
        "id": "eXi3WsdiFC86"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class VerletActionSpace:\n",
        "    def __init__(self, delta_interval=(-18.0, 18.0), num_bins=13, step_frequency=2.0):\n",
        "        self.delta_interval = delta_interval\n",
        "        self.num_bins = num_bins\n",
        "        self.step_frequency = step_frequency\n",
        "        self.time_step = 1.0 / step_frequency\n",
        "        self.accel_bin_edges = np.linspace(delta_interval[0], delta_interval[1], num_bins + 1)\n",
        "        self.accel_bin_centers = (self.accel_bin_edges[:-1] + self.accel_bin_edges[1:]) / 2\n",
        "\n",
        "        # Create the vocabulary as the Cartesian product of acceleration bins for x and y\n",
        "        self.vocab_size = num_bins * num_bins\n",
        "        self.token_to_accel = np.array([\n",
        "            (ax, ay) for ax in self.accel_bin_centers for ay in self.accel_bin_centers\n",
        "        ])  # Shape (169, 2)\n",
        "\n",
        "    def discretize_acceleration(self, accel):\n",
        "        # accel: array of shape (..., 2)\n",
        "        indices_x = np.digitize(accel[..., 0], self.accel_bin_edges) - 1\n",
        "        indices_y = np.digitize(accel[..., 1], self.accel_bin_edges) - 1\n",
        "        indices_x = np.clip(indices_x, 0, self.num_bins - 1)\n",
        "        indices_y = np.clip(indices_y, 0, self.num_bins - 1)\n",
        "        token_indices = indices_x * self.num_bins + indices_y\n",
        "        return token_indices  # Shape (...)\n",
        "\n",
        "    def continuous_acceleration(self, token_indices):\n",
        "        # token_indices: array of shape (...)\n",
        "        accel = self.token_to_accel[token_indices]\n",
        "        return accel  # Shape (..., 2)\n",
        "\n",
        "    def verlet_update(self, position, velocity, acceleration):\n",
        "        # position, velocity, acceleration: arrays of shape (..., 2)\n",
        "        new_position = position + velocity * self.time_step + 0.5 * acceleration * self.time_step ** 2\n",
        "        new_velocity = velocity + acceleration * self.time_step\n",
        "        return new_position, new_velocity\n",
        "\n",
        "def encode_trajectory(positions, velocities, action_space):\n",
        "    \"\"\"\n",
        "    Encode continuous trajectories into motion tokens.\n",
        "    positions: array of shape (T+1, 2)\n",
        "    velocities: array of shape (T+1, 2)\n",
        "    Returns:\n",
        "        token_indices: array of shape (T,)\n",
        "    \"\"\"\n",
        "    time_step = action_space.time_step\n",
        "    # Compute accelerations using inverse Verlet integration\n",
        "    accelerations = (velocities[1:] - velocities[:-1]) / time_step\n",
        "    # Discretize accelerations\n",
        "    token_indices = action_space.discretize_acceleration(accelerations)\n",
        "    return token_indices\n",
        "\n",
        "def decode_tokens(start_position, start_velocity, token_indices, action_space):\n",
        "    \"\"\"\n",
        "    Decode motion tokens back into continuous trajectories.\n",
        "    start_position: array of shape (2,)\n",
        "    start_velocity: array of shape (2,)\n",
        "    token_indices: array of shape (T,)\n",
        "    Returns:\n",
        "        positions: array of shape (T+1, 2)\n",
        "        velocities: array of shape (T+1, 2)\n",
        "    \"\"\"\n",
        "    positions = [start_position]\n",
        "    velocities = [start_velocity]\n",
        "    for idx in token_indices:\n",
        "        accel = action_space.continuous_acceleration(idx)\n",
        "        new_position, new_velocity = action_space.verlet_update(\n",
        "            positions[-1], velocities[-1], accel)\n",
        "        positions.append(new_position)\n",
        "        velocities.append(new_velocity)\n",
        "    positions = np.stack(positions)\n",
        "    velocities = np.stack(velocities)\n",
        "    return positions, velocities\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "TE8RqqQ2Q74w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Authenticate and setup\n",
        "auth.authenticate_user()\n",
        "\n",
        "# Configure dataset\n",
        "config = _config.DatasetConfig(\n",
        "    path='/content/training_tfexample.tfrecord',\n",
        "    data_format=_config.DataFormat.TFRECORD,\n",
        "    max_num_objects=32\n",
        ")\n",
        "\n",
        "# Create model config\n",
        "model_config = WayformerTrainingConfig()\n",
        "\n",
        "# Create training dataset\n",
        "print(\"Creating training dataset...\")\n",
        "action_space = VerletActionSpace()\n",
        "training_data = create_training_dataset(config, action_space, num_scenarios=1000)\n"
      ],
      "metadata": {
        "id": "7PSbd3VyJNSK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# Train model\n",
        "print(\"Training Wayformer model...\")\n",
        "model = train_wayformer(model_config, action_space, training_data)\n",
        "\n",
        "# Evaluate on a test scenario\n",
        "test_scenario = next(dataloader.simulator_state_generator(config=config))\n",
        "test_data = process_waymax_data(test_scenario)\n",
        "\n",
        "predictions = model({\n",
        "    'ego_in': test_data['agent_features'][:1, 0],\n",
        "    'agents_in': test_data['agent_features'][:1, 1:],\n",
        "    'roads': test_data['road_features'][:1]\n",
        "}, training=False)\n",
        "\n",
        "# Visualize predictions\n",
        "imgs = []\n",
        "state = test_scenario\n",
        "pred_trajectory = predictions['predicted_trajectory'][0]  # Take most likely trajectory\n",
        "\n",
        "for t in range(state.remaining_timesteps):\n",
        "    state = datatypes.update_state_by_log(state, num_steps=1)\n",
        "    img = visualization.plot_simulator_state(\n",
        "        state,\n",
        "        use_log_traj=True,\n",
        "        additional_trajectories={0: pred_trajectory[t]}  # Overlay predictions\n",
        "    )\n",
        "    imgs.append(img)\n",
        "\n",
        "mediapy.show_video(imgs, fps=10)"
      ],
      "metadata": {
        "id": "XOlvmsAsMxeB"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}