waymax_demo_with_model.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_with_model.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",
        "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": "code",
      "source": [
        "@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",
        "class Wayformer(tf.keras.Model):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "\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",
        "        self.max_points_per_lane = config['max_points_per_lane']\n",
        "        self.max_num_roads = config['max_num_roads']\n",
        "        self.num_queries_enc = config['num_queries_enc']\n",
        "        self.num_queries_dec = config['num_queries_dec']\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.agents_positional_embedding = self.add_weight(\n",
        "            shape=(1, 1, self._M + 1, self.d_k),\n",
        "            initializer='zeros',\n",
        "            trainable=True,\n",
        "            name='agents_pos_emb'\n",
        "        )\n",
        "\n",
        "        self.temporal_positional_embedding = self.add_weight(\n",
        "            shape=(1, self.past_T, 1, self.d_k),\n",
        "            initializer='zeros',\n",
        "            trainable=True,\n",
        "            name='temporal_pos_emb'\n",
        "        )\n",
        "\n",
        "        # Perceiver components\n",
        "        self.perceiver_encoder = PerceiverEncoder(\n",
        "            num_latents=self.num_queries_enc,\n",
        "            num_latent_channels=self.d_k\n",
        "        )\n",
        "\n",
        "        self.output_query_provider = TrainableQueryProvider(\n",
        "            num_queries=self.num_queries_dec,\n",
        "            num_query_channels=self.d_k,\n",
        "            init_scale=0.1\n",
        "        )\n",
        "\n",
        "        # Output heads\n",
        "        self.prob_predictor = tf.keras.layers.Dense(1)\n",
        "        self.output_model = tf.keras.layers.Dense(5 * self.T)\n",
        "\n",
        "        self.selu = tf.keras.layers.Activation('selu')\n",
        "\n",
        "    def call(self, inputs, training=False):\n",
        "        ego_in = inputs['ego_in']\n",
        "        agents_in = inputs['agents_in']\n",
        "        roads = inputs['roads']\n",
        "\n",
        "        # Process observations\n",
        "        ego_tensor, opps_tensor, opps_masks, env_masks = self.process_observations(ego_in, agents_in)\n",
        "\n",
        "        # Encode inputs\n",
        "        agents_tensor = tf.concat([tf.expand_dims(ego_tensor, 2), opps_tensor], axis=2)\n",
        "        agents_emb = self.selu(self.agents_dynamic_encoder(agents_tensor))\n",
        "\n",
        "        # Add positional embeddings\n",
        "        agents_emb = agents_emb + self.agents_positional_embedding + self.temporal_positional_embedding\n",
        "        batch_size = tf.shape(agents_emb)[0]\n",
        "        agents_emb = tf.reshape(agents_emb, [batch_size, -1, self.d_k])\n",
        "\n",
        "        # Process road features\n",
        "        road_pts_feats = self.selu(self.road_pts_lin(roads[:, :self.max_num_roads, :, :self.map_attr]))\n",
        "        road_pts_feats = tf.reshape(road_pts_feats, [batch_size, -1, self.d_k])\n",
        "\n",
        "        # Combine features\n",
        "        mixed_input_features = tf.concat([agents_emb, road_pts_feats], axis=1)\n",
        "\n",
        "        # Encode with Perceiver\n",
        "        context = self.perceiver_encoder(mixed_input_features, training=training)\n",
        "\n",
        "        # Generate outputs\n",
        "        output_query = self.output_query_provider()\n",
        "        out_seq = tf.matmul(output_query, context, transpose_b=True)\n",
        "\n",
        "        # Predict trajectories and probabilities\n",
        "        out_dists = self.output_model(out_seq[:, :self.c])\n",
        "        out_dists = tf.reshape(out_dists, [batch_size, self.c, self.T, -1])\n",
        "\n",
        "        mode_probs = self.prob_predictor(out_seq[:, :self.c])\n",
        "        mode_probs = tf.reshape(mode_probs, [batch_size, self.c])\n",
        "\n",
        "        return {\n",
        "            'predicted_probability': mode_probs,\n",
        "            'predicted_trajectory': out_dists,\n",
        "            'scene_emb': tf.reshape(out_seq[:, :self.num_queries_dec], [batch_size, -1])\n",
        "        }\n",
        "\n",
        "    def process_observations(self, ego, agents):\n",
        "        # Process ego vehicle data\n",
        "        ego_tensor = ego[:, :, :self.k_attr]\n",
        "        env_masks_orig = ego[:, :, -1]\n",
        "        env_masks = tf.cast(1.0 - env_masks_orig, tf.bool)\n",
        "        env_masks = tf.expand_dims(env_masks, 1)\n",
        "        env_masks = tf.tile(env_masks, [1, self.num_queries_dec, 1])\n",
        "        batch_size = tf.shape(ego)[0]\n",
        "        env_masks = tf.reshape(env_masks, [batch_size * self.num_queries_dec, -1])\n",
        "\n",
        "        # Process other agents data\n",
        "        temp_masks = tf.concat([\n",
        "            tf.ones_like(tf.expand_dims(env_masks_orig, -1)),\n",
        "            agents[:, :, :, -1]\n",
        "        ], axis=-1)\n",
        "        opps_masks = tf.cast(1.0 - temp_masks, tf.bool)\n",
        "        opps_tensor = agents[:, :, :, :self.k_attr]\n",
        "\n",
        "        return ego_tensor, opps_tensor, opps_masks, env_masks\n",
        "\n",
        "    def train_step(self, data):\n",
        "        with tf.GradientTape() as tape:\n",
        "            predictions = self(data['input_dict'], training=True)\n",
        "            loss = self.compute_loss(predictions, data)\n",
        "\n",
        "        gradients = tape.gradient(loss, self.trainable_variables)\n",
        "        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))\n",
        "\n",
        "        return {'loss': loss}\n",
        "\n",
        "    def compute_loss(self, predictions, data):\n",
        "        # Implement your loss computation here\n",
        "        # This should match the PyTorch Criterion class functionality\n",
        "        pass\n"
      ],
      "metadata": {
        "id": "YX5nfz3CHxL1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class WayformerTrainingConfig:\n",
        "    def __init__(self):\n",
        "        self.num_map_feature = 11  # Road feature dimensions\n",
        "        self.num_agent_feature = 9  # Agent feature dimensions\n",
        "        self.hidden_size = 256\n",
        "        self.max_num_agents = 32\n",
        "        self.num_modes = 6\n",
        "        self.future_len = 80  # 8 seconds with 10Hz\n",
        "        self.past_len = 11   # 1 second with 10Hz\n",
        "        self.dropout = 0.1\n",
        "        self.tx_num_heads = 8\n",
        "        self.max_points_per_lane = 40\n",
        "        self.max_num_roads = 50\n",
        "        self.num_queries_enc = 128\n",
        "        self.num_queries_dec = 64\n",
        "        self.learning_rate = 1e-4\n",
        "        self.batch_size = 32\n",
        "        self.num_epochs = 10\n",
        "\n",
        "def process_waymax_data(state):\n",
        "    \"\"\"Convert waymax state to model input format\"\"\"\n",
        "    # Extract relevant features\n",
        "    positions = state.object_metadata.position\n",
        "    velocities = state.object_metadata.velocity\n",
        "    headings = state.object_metadata.heading\n",
        "    valid_mask = state.object_metadata.valid\n",
        "\n",
        "    # Combine features for agents\n",
        "    agent_features = np.concatenate([\n",
        "        positions,                     # x, y positions\n",
        "        velocities,                    # vx, vy velocities\n",
        "        np.expand_dims(headings, -1),  # heading angle\n",
        "        np.zeros_like(positions[..., :3]),  # placeholder for additional features\n",
        "        np.expand_dims(valid_mask, -1) # validity mask\n",
        "    ], axis=-1)\n",
        "\n",
        "    # Extract road features from state\n",
        "    road_features = state.roadgraph_segments_pos\n",
        "    road_valid = state.roadgraph_segments_valid\n",
        "\n",
        "    return {\n",
        "        'agent_features': agent_features,\n",
        "        'road_features': road_features,\n",
        "        'road_valid': road_valid,\n",
        "        'target_features': state.object_trajectories  # Future trajectories\n",
        "    }\n",
        "\n",
        "def create_training_dataset(config, num_scenarios=1000):\n",
        "    \"\"\"Create training dataset from waymax data\"\"\"\n",
        "    data_iter = dataloader.simulator_state_generator(config=config)\n",
        "    dataset = []\n",
        "\n",
        "    for _ in tqdm(range(num_scenarios)):\n",
        "        try:\n",
        "            scenario = next(data_iter)\n",
        "            processed_data = process_waymax_data(scenario)\n",
        "            dataset.append(processed_data)\n",
        "        except StopIteration:\n",
        "            break\n",
        "\n",
        "    return dataset\n",
        "\n",
        "def train_wayformer(model_config, training_data):\n",
        "    \"\"\"Train Wayformer model on processed data\"\"\"\n",
        "    # Initialize model\n",
        "    model = Wayformer(model_config)\n",
        "    optimizer = tf.keras.optimizers.Adam(learning_rate=model_config.learning_rate)\n",
        "\n",
        "    # Training loop\n",
        "    num_batches = len(training_data) // model_config.batch_size\n",
        "\n",
        "    for epoch in range(model_config.num_epochs):\n",
        "        total_loss = 0\n",
        "\n",
        "        for batch_idx in tqdm(range(num_batches)):\n",
        "            batch_data = training_data[batch_idx * model_config.batch_size:\n",
        "                                     (batch_idx + 1) * model_config.batch_size]\n",
        "\n",
        "            with tf.GradientTape() as tape:\n",
        "                # Forward pass\n",
        "                predictions = model({\n",
        "                    'ego_in': batch_data['agent_features'][:, 0],  # First agent as ego\n",
        "                    'agents_in': batch_data['agent_features'][:, 1:],  # Other agents\n",
        "                    'roads': batch_data['road_features']\n",
        "                }, training=True)\n",
        "\n",
        "                # Compute loss\n",
        "                loss = compute_trajectory_loss(\n",
        "                    predictions['predicted_trajectory'],\n",
        "                    predictions['predicted_probability'],\n",
        "                    batch_data['target_features']\n",
        "                )\n",
        "\n",
        "            # Backpropagation\n",
        "            gradients = tape.gradient(loss, model.trainable_variables)\n",
        "            optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "\n",
        "            total_loss += loss\n",
        "\n",
        "        avg_loss = total_loss / num_batches\n",
        "        print(f\"Epoch {epoch + 1}/{model_config.num_epochs}, Average Loss: {avg_loss:.4f}\")\n",
        "\n",
        "    return model\n",
        "\n",
        "def compute_trajectory_loss(pred_trajectories, pred_probabilities, target_trajectories):\n",
        "    \"\"\"Compute training loss\"\"\"\n",
        "    # Implement multi-modal trajectory prediction loss\n",
        "    displacement_error = tf.reduce_mean(tf.sqrt(tf.reduce_sum(\n",
        "        tf.square(pred_trajectories[..., :2] - target_trajectories[..., :2]), axis=-1)))\n",
        "\n",
        "    probability_loss = tf.reduce_mean(tf.keras.losses.categorical_crossentropy(\n",
        "        tf.ones_like(pred_probabilities) / pred_probabilities.shape[-1],\n",
        "        pred_probabilities))\n",
        "\n",
        "    return displacement_error + 0.1 * probability_loss"
      ],
      "metadata": {
        "id": "14cw1PyyH4SF"
      },
      "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",
        "training_data = create_training_dataset(config, num_scenarios=1000)\n",
        "\n",
        "# Train model\n",
        "print(\"Training Wayformer model...\")\n",
        "model = train_wayformer(model_config, 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": "xyvX9GGmI_FV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "7PSbd3VyJNSK"
      },
      "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
}