waymax_demo_kvcache.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_kvcache.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": "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 ModuleOutput:\n",
        "    last_hidden_state: tf.Tensor\n",
        "    kv_cache: Optional[Tuple[tf.Tensor, tf.Tensor]] = None\n"
      ],
      "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",
        "    num_timesteps = trajectory.x.shape[1]\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, 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 prepare_model_input_with_history(batch_data, history_motion_tokens, T=80):\n",
        "    \"\"\"Prepare batch data for model input with history buffers and 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",
        "        '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",
        "        # We will not use 'future_features' and 'motion_tokens' from batch_data\n",
        "        # as we are replacing them with the history buffers\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",
        "    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 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 = agents[:31]  # Limit to 31 other agents\n",
        "\n",
        "        # Pad if we have fewer than 31 agents\n",
        "        if len(agents) < 31:\n",
        "            pad_length = 31 - len(agents)\n",
        "            # Create padding agents with zeros\n",
        "            agent_pad = jnp.zeros((pad_length, ego_feat.shape[0], ego_feat.shape[1]))\n",
        "            agents.extend([agent_pad])\n",
        "\n",
        "        agents = jnp.stack(agents)\n",
        "        other_agents.append(agents)\n",
        "\n",
        "    # Prepare history buffers for motion tokens\n",
        "    # Ensure history buffers have length T and proper padding\n",
        "\n",
        "    # For motion tokens\n",
        "    if len(history_motion_tokens) < T:\n",
        "        pad_length = T - len(history_motion_tokens)\n",
        "        # Pad with zeros or appropriate padding tokens\n",
        "        padding_motion_tokens = [0] * pad_length  # Assuming 0 is the padding token\n",
        "        history_motion_tokens = padding_motion_tokens + history_motion_tokens\n",
        "    else:\n",
        "        history_motion_tokens = history_motion_tokens[-T:]  # Keep the last T tokens\n",
        "\n",
        "    # Convert to numpy array and add batch dimension\n",
        "    motion_tokens_array = np.array(history_motion_tokens)[np.newaxis, :]  # Shape: (1, T)\n",
        "\n",
        "    # Stack the features\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",
        "        'motion_tokens': motion_tokens_array        # (batch_size, T)\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": [
        "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": [
        "@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 = 1024  # 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"
      ],
      "metadata": {
        "id": "DhnY7_OWHWLs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiHeadAttentionCache(tf.keras.layers.Layer):\n",
        "    def __init__(\n",
        "        self,\n",
        "        num_heads: int,\n",
        "        d_model: int,\n",
        "        dropout_rate: float = 0.1\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.num_heads = num_heads\n",
        "        self.d_model = d_model\n",
        "        self.depth = d_model // num_heads\n",
        "\n",
        "        assert d_model % num_heads == 0, \"d_model must be divisible by num_heads\"\n",
        "\n",
        "        self.wq = tf.keras.layers.Dense(d_model)\n",
        "        self.wk = tf.keras.layers.Dense(d_model)\n",
        "        self.wv = tf.keras.layers.Dense(d_model)\n",
        "\n",
        "        self.dense = tf.keras.layers.Dense(d_model)\n",
        "        self.dropout = tf.keras.layers.Dropout(dropout_rate)\n",
        "\n",
        "    def split_heads(self, x, batch_size):\n",
        "        \"\"\"Split the last dimension into (num_heads, depth).\n",
        "        Transpose the result to shape (batch_size, num_heads, seq_len, depth)\"\"\"\n",
        "        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
        "        return tf.transpose(x, perm=[0, 2, 1, 3])\n",
        "\n",
        "    def call(self, v, k, q, mask, cache=None, cache_key=None, training=False):\n",
        "        batch_size = tf.shape(q)[0]\n",
        "\n",
        "        q = self.wq(q)  # (batch_size, seq_len_q, d_model)\n",
        "        k = self.wk(k)  # (batch_size, seq_len_k, d_model)\n",
        "        v = self.wv(v)  # (batch_size, seq_len_v, d_model)\n",
        "\n",
        "        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n",
        "        k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n",
        "        v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n",
        "\n",
        "        if cache is not None and cache_key is not None:\n",
        "            # Use the cache instance\n",
        "            cache_instance = cache[cache_key]\n",
        "            # Update the cache with only the new keys and values\n",
        "            cache_instance.update(k, v)\n",
        "            # Retrieve the cached keys and values\n",
        "            k_cache, v_cache = cache_instance.get_key_value()\n",
        "        else:\n",
        "            k_cache = k\n",
        "            v_cache = v\n",
        "        #print shape of k_cache and v_cache\n",
        "        # print(f\"k_cache shape: {k_cache.shape}\")\n",
        "        # print(f\"v_cache shape: {v_cache.shape}\")\n",
        "        # Scaled dot-product attention\n",
        "        matmul_qk = tf.matmul(q, k_cache, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
        "        dk = tf.cast(tf.shape(k_cache)[-1], tf.float32)\n",
        "        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
        "\n",
        "        # Apply mask if provided\n",
        "        if mask is not None:\n",
        "            scaled_attention_logits += (mask * -1e9)\n",
        "\n",
        "        # Softmax on the last axis (seq_len_k)\n",
        "        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)\n",
        "        attention_weights = self.dropout(attention_weights, training=training)\n",
        "\n",
        "        output = tf.matmul(attention_weights, v_cache)  # (..., seq_len_q, depth)\n",
        "\n",
        "        output = tf.transpose(output, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)\n",
        "        concat_output = tf.reshape(output, (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)\n",
        "\n",
        "        output = self.dense(concat_output)  # (batch_size, seq_len_q, d_model)\n",
        "\n",
        "        return output, attention_weights\n",
        "\n",
        "class TransformerDecoderLayerCache(tf.keras.layers.Layer):\n",
        "    def __init__(self, d_model, num_heads, dff, dropout_rate=0.1):\n",
        "        super(TransformerDecoderLayerCache, self).__init__()\n",
        "\n",
        "        # Self-attention (masked)\n",
        "        self.mha1 = MultiHeadAttentionCache(num_heads, d_model, dropout_rate)\n",
        "\n",
        "        # Cross-attention with encoder outputs\n",
        "        self.mha2 = MultiHeadAttentionCache(num_heads, d_model, 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, cache=None, layer_idx=0):\n",
        "        # Self-attention (masked)\n",
        "        attn1, _ = self.mha1(\n",
        "            q=x,\n",
        "            k=x,\n",
        "            v=x,\n",
        "            mask=look_ahead_mask,\n",
        "            cache=cache,\n",
        "            cache_key=f'decoder_layer_{layer_idx}_self_attn',\n",
        "            training=training\n",
        "        )\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",
        "            q=out1,\n",
        "            k=enc_output,\n",
        "            v=enc_output,\n",
        "            mask=padding_mask,\n",
        "            training=training\n",
        "        )\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)\n",
        "        ffn_output = self.dropout3(ffn_output, training=training)\n",
        "        out3 = self.layernorm3(out2 + ffn_output)\n",
        "\n",
        "        return out3\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 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",
        "        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 = MultiHeadAttentionCache(\n",
        "            num_heads=num_cross_attention_heads,\n",
        "            d_model=num_latent_channels,\n",
        "            dropout_rate=dropout\n",
        "        )\n",
        "\n",
        "        self.self_attention = MultiHeadAttentionCache(\n",
        "            num_heads=num_cross_attention_heads,\n",
        "            d_model=num_latent_channels,\n",
        "            dropout_rate=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(\n",
        "            v=x,\n",
        "            k=x,\n",
        "            q=x_latent,\n",
        "            mask=pad_mask,\n",
        "            training=training\n",
        "        )\n",
        "        x_latent = residual + self.dropout(cross_attn_output, 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(\n",
        "            v=x_latent,\n",
        "            k=x_latent,\n",
        "            q=x_latent,\n",
        "            mask=None,\n",
        "            training=training\n",
        "        )\n",
        "        x_latent = residual + self.dropout(self_attn_output, training=training)\n",
        "\n",
        "        return x_latent\n",
        "\n",
        "class WayformerCache(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=(self.T, 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",
        "            TransformerDecoderLayerCache(\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",
        "        self.built = True\n",
        "\n",
        "    def encode_agents(self, ego_in, agents_in, training=False):\n",
        "        \"\"\"Process agent features\"\"\"\n",
        "        batch_size = tf.shape(ego_in)[0]\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",
        "\n",
        "        # Encode agent features\n",
        "        encoded = self.agent_encoder(all_agents, training=training)\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",
        "        # Apply self-attention\n",
        "        attended = self.agent_transformer(\n",
        "            query=encoded,\n",
        "            key=encoded,\n",
        "            value=encoded,\n",
        "            training=training\n",
        "        )\n",
        "\n",
        "        # Final reshape\n",
        "        final_encoded = tf.reshape(attended, [batch_size, -1, self.d_k])\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, cache=None, training=False):\n",
        "        \"\"\"Forward pass with KV cache support\"\"\"\n",
        "        ego_in = inputs['ego_in']\n",
        "        agents_in = inputs['agents_in']\n",
        "        roads = inputs['roads']\n",
        "        motion_tokens = inputs['motion_tokens']  # Shape: (batch_size, T)\n",
        "\n",
        "        # Process agents\n",
        "        agents_encoded = self.encode_agents(ego_in, agents_in, training)\n",
        "        # Process roads\n",
        "        roads_encoded = self.encode_roads(roads, training)\n",
        "        # Combine features\n",
        "        combined_features = tf.concat([agents_encoded, roads_encoded], axis=1)\n",
        "        # Apply perceiver encoding\n",
        "        context = self.perceiver(combined_features, training=training)\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",
        "\n",
        "        # Get trajectory predictions\n",
        "        output_features = tf.matmul(query, context, transpose_b=True)\n",
        "        trajectories = self.trajectory_projection(output_features)\n",
        "        trajectories = tf.reshape(trajectories, [batch_size, self.c, self.T, -1])\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",
        "\n",
        "        # Embed motion tokens\n",
        "        token_embeddings = self.token_embedding(motion_tokens)  # Shape: (batch_size, T, d_k)\n",
        "\n",
        "        # Add positional encoding\n",
        "        positional_encoding = tf.expand_dims(self.positional_encoding[:tf.shape(motion_tokens)[1], :], axis=0)\n",
        "        token_embeddings += positional_encoding  # Broadcasting over batch_size\n",
        "\n",
        "        # Prepare attention masks\n",
        "        batch_size = tf.shape(token_embeddings)[0]\n",
        "        seq_len = tf.shape(token_embeddings)[1]\n",
        "        look_ahead_mask = self.create_look_ahead_mask(seq_len)  # Shape: (1, seq_len, seq_len)\n",
        "\n",
        "        # Pass through decoder layers\n",
        "        decoder_output = token_embeddings\n",
        "        for idx, layer in enumerate(self.decoder_layers):\n",
        "            decoder_output = layer(\n",
        "                decoder_output,\n",
        "                context,\n",
        "                training=training,\n",
        "                look_ahead_mask=look_ahead_mask,\n",
        "                cache=cache,\n",
        "                layer_idx=idx\n",
        "            )\n",
        "\n",
        "        # Predict logits for next tokens\n",
        "        logits = self.token_projection(decoder_output)  # Shape: (batch_size, seq_len, 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",
        "        }\n",
        "\n",
        "    def create_look_ahead_mask(self, size):\n",
        "        mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
        "        mask = tf.cast(mask, tf.float32)  # Shape: (size, size)\n",
        "        return mask[tf.newaxis, :, :]  # Shape: (1, size, size)"
      ],
      "metadata": {
        "id": "6U8ohi05z_3x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\" This cache is a baseline KV cache, storing all past keys and vals\"\"\"\n",
        "from abc import ABC, abstractmethod\n",
        "\n",
        "class BaseCache(ABC):\n",
        "    @abstractmethod\n",
        "    def update(self, key, value):\n",
        "        \"\"\"Update the cache with new key and value.\"\"\"\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def get_key_value(self):\n",
        "        \"\"\"Retrieve the current key and value from the cache.\"\"\"\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def reset(self):\n",
        "        \"\"\"Reset the cache to its initial state.\"\"\"\n",
        "        pass\n",
        "\n",
        "class FullCache(BaseCache):\n",
        "    def __init__(self, max_cache_length, num_heads, depth):\n",
        "        self.max_cache_length = max_cache_length\n",
        "        self.num_heads = num_heads\n",
        "        self.depth = depth\n",
        "        self.current_length = 0\n",
        "\n",
        "        self.key = None\n",
        "        self.value = None\n",
        "\n",
        "    def update(self, key, value):\n",
        "        batch_size = tf.shape(key)[0]\n",
        "        seq_len = tf.shape(key)[2]  # seq_len_k\n",
        "\n",
        "        if self.key is None:\n",
        "            # Initialize the variables with maximum cache length\n",
        "            self.key = tf.Variable(\n",
        "                tf.zeros([batch_size, self.num_heads, self.max_cache_length, self.depth], dtype=key.dtype),\n",
        "                trainable=False)\n",
        "            self.value = tf.Variable(\n",
        "                tf.zeros([batch_size, self.num_heads, self.max_cache_length, self.depth], dtype=value.dtype),\n",
        "                trainable=False)\n",
        "            self.current_length = 0\n",
        "\n",
        "        if self.current_length + seq_len > self.max_cache_length:\n",
        "            raise ValueError(\"Cache overflow: Exceeded maximum cache length\")\n",
        "\n",
        "        # Assign the new key and value to the appropriate slice\n",
        "        self.key[:, :, self.current_length:self.current_length + seq_len, :].assign(key)\n",
        "        self.value[:, :, self.current_length:self.current_length + seq_len, :].assign(value)\n",
        "\n",
        "        self.current_length += seq_len\n",
        "\n",
        "    def get_key_value(self):\n",
        "        if self.key is None or self.current_length == 0:\n",
        "            return None, None\n",
        "        # Return the key and value up to current_length\n",
        "        key = self.key[:, :, :self.current_length, :]\n",
        "        value = self.value[:, :, :self.current_length, :]\n",
        "        return key, value\n",
        "\n",
        "    def reset(self):\n",
        "        self.current_length = 0\n",
        "        if self.key is not None:\n",
        "            self.key.assign(tf.zeros_like(self.key))\n",
        "            self.value.assign(tf.zeros_like(self.value))\n",
        "\n",
        "class SlidingWindowCache(BaseCache):\n",
        "    def __init__(self, window_size, num_heads, head_dim):\n",
        "        \"\"\"\n",
        "        Initializes the SlidingWindowCache with a fixed window size.\n",
        "\n",
        "        Args:\n",
        "            window_size (int): Fixed size of the sliding window.\n",
        "            num_heads (int): Number of attention heads.\n",
        "            head_dim (int): Dimension per attention head.\n",
        "        \"\"\"\n",
        "        self.window_size = window_size\n",
        "        self.num_heads = num_heads\n",
        "        self.head_dim = head_dim\n",
        "\n",
        "        # Initialize pointers for the ring buffer\n",
        "        self.current_size = tf.Variable(0, trainable=False, dtype=tf.int32)\n",
        "        self.start_idx = tf.Variable(0, trainable=False, dtype=tf.int32)\n",
        "        self.end_idx = tf.Variable(0, trainable=False, dtype=tf.int32)\n",
        "\n",
        "        # Preallocate buffers for keys and values with fixed size\n",
        "        self.key_buffer = None\n",
        "        self.value_buffer = None\n",
        "\n",
        "    def update(self, key, value):\n",
        "        \"\"\"\n",
        "        Updates the cache with new key and value tensors.\n",
        "\n",
        "        Args:\n",
        "            key (tf.Tensor): Tensor of shape [batch_size, num_heads, seq_len_k, head_dim].\n",
        "            value (tf.Tensor): Tensor of shape [batch_size, num_heads, seq_len_v, head_dim].\n",
        "        \"\"\"\n",
        "        batch_size = tf.shape(key)[0]\n",
        "        seq_len_kv = tf.shape(key)[2]  # Typically 1 during autoregressive inference\n",
        "\n",
        "        # Enforce that seq_len_kv=1\n",
        "        if seq_len_kv != 1:\n",
        "            raise ValueError(f\"SlidingWindowCache only supports seq_len_kv=1, but got {seq_len_kv}\")\n",
        "\n",
        "        # Initialize buffers upon the first update\n",
        "        if self.key_buffer is None:\n",
        "            self.key_buffer = tf.Variable(\n",
        "                tf.zeros([batch_size, self.num_heads, self.window_size, self.head_dim], dtype=key.dtype),\n",
        "                trainable=False\n",
        "            )\n",
        "            self.value_buffer = tf.Variable(\n",
        "                tf.zeros([batch_size, self.num_heads, self.window_size, self.head_dim], dtype=value.dtype),\n",
        "                trainable=False\n",
        "            )\n",
        "            self.current_size.assign(0)\n",
        "            self.start_idx.assign(0)\n",
        "            self.end_idx.assign(0)\n",
        "\n",
        "        # Check if adding a new token exceeds window_size\n",
        "        if self.current_size.numpy() >= self.window_size:\n",
        "            # Overwrite the oldest token\n",
        "            self.key_buffer[:, :, self.start_idx, :].assign(key[:, :, 0, :])\n",
        "            self.value_buffer[:, :, self.start_idx, :].assign(value[:, :, 0, :])\n",
        "\n",
        "            # Move the window forward\n",
        "            self.start_idx.assign((self.start_idx + 1) % self.window_size)\n",
        "            self.end_idx.assign((self.end_idx + 1) % self.window_size)\n",
        "        else:\n",
        "            # Add the new token at end_idx\n",
        "            self.key_buffer[:, :, self.end_idx, :].assign(key[:, :, 0, :])\n",
        "            self.value_buffer[:, :, self.end_idx, :].assign(value[:, :, 0, :])\n",
        "\n",
        "            # Move the end_idx forward\n",
        "            self.end_idx.assign((self.end_idx + 1) % self.window_size)\n",
        "            self.current_size.assign(self.current_size + 1)\n",
        "\n",
        "    def get_key_value(self):\n",
        "        \"\"\"\n",
        "        Retrieves the concatenated key and value tensors up to the current cache length.\n",
        "\n",
        "        Returns:\n",
        "            Tuple[tf.Tensor, tf.Tensor]: Concatenated key and value tensors.\n",
        "        \"\"\"\n",
        "        if self.key_buffer is None or self.current_size.numpy() == 0:\n",
        "            return None, None\n",
        "\n",
        "        if self.start_idx.numpy() < self.end_idx.numpy():\n",
        "            # Continuous slice\n",
        "            key = self.key_buffer[:, :, self.start_idx:self.end_idx, :]\n",
        "            value = self.value_buffer[:, :, self.start_idx:self.end_idx, :]\n",
        "        else:\n",
        "            # Wrap-around slice\n",
        "            key_part1 = self.key_buffer[:, :, self.start_idx:, :]\n",
        "            key_part2 = self.key_buffer[:, :, :self.end_idx, :]\n",
        "            key = tf.concat([key_part1, key_part2], axis=2)\n",
        "\n",
        "            value_part1 = self.value_buffer[:, :, self.start_idx:, :]\n",
        "            value_part2 = self.value_buffer[:, :, :self.end_idx, :]\n",
        "            value = tf.concat([value_part1, value_part2], axis=2)\n",
        "\n",
        "        return key, value\n",
        "\n",
        "    def reset(self):\n",
        "        \"\"\"\n",
        "        Resets the cache by clearing all key and value data.\n",
        "        \"\"\"\n",
        "        self.current_size.assign(0)\n",
        "        self.start_idx.assign(0)\n",
        "        self.end_idx.assign(0)\n",
        "        if self.key_buffer is not None:\n",
        "            self.key_buffer.assign(tf.zeros_like(self.key_buffer))\n",
        "            self.value_buffer.assign(tf.zeros_like(self.value_buffer))"
      ],
      "metadata": {
        "id": "ChV5IHfQaUIt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def pad_state_to_num_timesteps(state, required_timesteps):\n",
        "    \"\"\"\n",
        "    Pads the state's sim_trajectory arrays to have at least required_timesteps.\n",
        "    Returns the updated state with JAX arrays.\n",
        "    \"\"\"\n",
        "    import jax.numpy as jnp\n",
        "\n",
        "    # Get the existing number of timesteps\n",
        "    existing_timesteps = state.sim_trajectory.x.shape[1]\n",
        "    timesteps_to_add = required_timesteps - existing_timesteps\n",
        "\n",
        "    if timesteps_to_add <= 0:\n",
        "        # No padding needed\n",
        "        return state\n",
        "\n",
        "    # Function to pad an array along the second dimension (timesteps)\n",
        "    def pad_array(array, pad_length, pad_value=0):\n",
        "        pad_shape = list(array.shape)\n",
        "        pad_shape[1] = pad_length\n",
        "        if array.dtype == jnp.bool_:\n",
        "            padding = jnp.full(pad_shape, pad_value, dtype=array.dtype)\n",
        "        else:\n",
        "            padding = jnp.full(pad_shape, pad_value, dtype=array.dtype)\n",
        "        return jnp.concatenate([array, padding], axis=1)\n",
        "\n",
        "    # Pad each relevant array in sim_trajectory\n",
        "    padded_sim_trajectory = state.sim_trajectory.replace(\n",
        "        x=pad_array(state.sim_trajectory.x, timesteps_to_add),\n",
        "        y=pad_array(state.sim_trajectory.y, timesteps_to_add),\n",
        "        z=pad_array(state.sim_trajectory.z, timesteps_to_add),\n",
        "        vel_x=pad_array(state.sim_trajectory.vel_x, timesteps_to_add),\n",
        "        vel_y=pad_array(state.sim_trajectory.vel_y, timesteps_to_add),\n",
        "        yaw=pad_array(state.sim_trajectory.yaw, timesteps_to_add),\n",
        "        length=pad_array(state.sim_trajectory.length, timesteps_to_add),\n",
        "        width=pad_array(state.sim_trajectory.width, timesteps_to_add),\n",
        "        height=pad_array(state.sim_trajectory.height, timesteps_to_add),\n",
        "        timestamp_micros=pad_array(\n",
        "            state.sim_trajectory.timestamp_micros, timesteps_to_add, pad_value=0\n",
        "        ),\n",
        "        valid=pad_array(state.sim_trajectory.valid, timesteps_to_add, pad_value=False)\n",
        "    )\n",
        "\n",
        "    # Return the updated state with the padded sim_trajectory\n",
        "    return state.replace(sim_trajectory=padded_sim_trajectory)\n",
        "\n",
        "import sys\n",
        "\n",
        "def get_cache_memory_usage(cache):\n",
        "    total_size = 0\n",
        "    for layer_key, cache_instance in cache.items():\n",
        "        k_cache, v_cache = cache_instance.get_key_value()\n",
        "        if k_cache is not None and v_cache is not None:\n",
        "            size_k = tf.size(k_cache)\n",
        "            size_v = tf.size(v_cache)\n",
        "            dtype_size = k_cache.dtype.size\n",
        "            layer_memory = (size_k + size_v) * dtype_size  # Bytes\n",
        "            total_size += layer_memory\n",
        "            layer_memory_mb = layer_memory / (1024 ** 2)\n",
        "            print(f\"Layer {layer_key} cache memory usage: {layer_memory_mb:.2f} MB\")\n",
        "    return total_size"
      ],
      "metadata": {
        "id": "gaKiy8OqJBYr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\"This block implements inference without KV cache\"\"\"\n",
        "import time\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "# Initialize the model\n",
        "model_config = WayformerTrainingConfig()\n",
        "action_space = VerletActionSpace()\n",
        "model = WayformerCache(model_config, action_space)\n",
        "\n",
        "# Get a test scenario\n",
        "test_scenario = next(dataloader.simulator_state_generator(config=config))\n",
        "\n",
        "# Number of autoregressive steps for inference benchmarking\n",
        "num_autoregressive_steps = 1024  # Or any desired number of steps\n",
        "\n",
        "# Initialize the state with the test scenario\n",
        "state = test_scenario\n",
        "\n",
        "# Pad the state to have sufficient timesteps\n",
        "initial_timestep = state.timestep\n",
        "required_timesteps = initial_timestep + num_autoregressive_steps + 1  # +1 for zero-based indexing\n",
        "state = pad_state_to_num_timesteps(state, required_timesteps)\n",
        "\n",
        "# Initialize the history buffer for motion tokens\n",
        "history_motion_tokens = []\n",
        "\n",
        "total_time = 0.0\n",
        "\n",
        "# Initialize the KV cache\n",
        "# If cache is None, initialize it\n",
        "cache = None\n",
        "elapsed_time1 = 0.0\n",
        "for step in range(num_autoregressive_steps):\n",
        "    # Process the current state to get model input\n",
        "    test_data = process_waymax_data(state, action_space)\n",
        "\n",
        "    # Update history buffers\n",
        "    if step == 0:\n",
        "        # For the first step, initialize with zeros or appropriate start tokens\n",
        "        initial_motion_token = 0  # Replace with your start token if applicable\n",
        "        history_motion_tokens.append(initial_motion_token)\n",
        "    else:\n",
        "        start_time = time.time()\n",
        "        # Append the predicted motion token from the previous step\n",
        "        history_motion_tokens.append(predicted_motion_token)\n",
        "        end_time = time.time()\n",
        "        elapsed_time1 = end_time - start_time\n",
        "        total_time += elapsed_time1\n",
        "\n",
        "    # Prepare the model inputs, including the history\n",
        "    # With KV cache, we only need the last token\n",
        "    model_input = prepare_model_input_with_history(\n",
        "        [test_data],\n",
        "        history_motion_tokens[-1:],  # Use only the last token\n",
        "        T=len(history_motion_tokens)\n",
        "    )\n",
        "\n",
        "    # Convert model inputs to tensors\n",
        "    model_inputs = {k: tf.convert_to_tensor(v, dtype=tf.float32) for k, v in model_input.items()}\n",
        "\n",
        "    # Perform inference\n",
        "    # Measure time before inference\n",
        "    start_time = time.time()\n",
        "\n",
        "    # Pass the cache to the model\n",
        "    predictions = model(model_inputs, cache=cache, training=False)\n",
        "\n",
        "    # Measure time after inference\n",
        "    end_time = time.time()\n",
        "    # Calculate elapsed time\n",
        "    elapsed_time2 = end_time - start_time\n",
        "    total_time += elapsed_time2\n",
        "    # Print the timing result\n",
        "    print(f\"Auto-regressive inference time for step {step}: {elapsed_time2+elapsed_time1:.6f} seconds\")\n",
        "\n",
        "    # Get the predicted motion token from logits\n",
        "    # 'logits' has shape (batch_size, 1, vocab_size)\n",
        "    # We need the predicted token for the current step\n",
        "    predicted_motion_token = tf.argmax(predictions['logits'], axis=-1).numpy()[0, 0]\n",
        "\n",
        "    # Convert motion token to acceleration\n",
        "    acceleration = action_space.continuous_acceleration(predicted_motion_token)  # Shape: (2,)\n",
        "\n",
        "    # Update the position and velocity using Verlet integration\n",
        "    ego_index = 0  # Assuming ego vehicle is at index 0\n",
        "\n",
        "    # Get current position and velocity from the state\n",
        "    current_position = np.array([\n",
        "        state.sim_trajectory.x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    current_velocity = np.array([\n",
        "        state.sim_trajectory.vel_x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.vel_y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    # Use the verlet_update function to compute the next position and velocity\n",
        "    new_position, new_velocity = action_space.verlet_update(current_position, current_velocity, acceleration)\n",
        "\n",
        "    # Optionally compute the new yaw angle based on the velocity vector\n",
        "    next_yaw = np.arctan2(new_velocity[1], new_velocity[0])  # Calculate yaw from velocity components\n",
        "\n",
        "    # Update the state with the predicted values\n",
        "    state = state.replace(\n",
        "        sim_trajectory=state.sim_trajectory.replace(\n",
        "            x=state.sim_trajectory.x.at[ego_index, state.timestep + 1].set(new_position[0]),\n",
        "            y=state.sim_trajectory.y.at[ego_index, state.timestep + 1].set(new_position[1]),\n",
        "            vel_x=state.sim_trajectory.vel_x.at[ego_index, state.timestep + 1].set(new_velocity[0]),\n",
        "            vel_y=state.sim_trajectory.vel_y.at[ego_index, state.timestep + 1].set(new_velocity[1]),\n",
        "            yaw=state.sim_trajectory.yaw.at[ego_index, state.timestep + 1].set(next_yaw),\n",
        "            valid=state.sim_trajectory.valid.at[ego_index, state.timestep + 1].set(True)\n",
        "        ),\n",
        "        timestep=state.timestep + 1  # Advance the timestep\n",
        "    )\n",
        "\n",
        "    # Break the loop if we've reached the end of available timesteps\n",
        "    if state.timestep >= state.sim_trajectory.x.shape[1] - 1:\n",
        "        break\n",
        "\n",
        "# Print total time\n",
        "print(f\"Total auto-regressive inference time with KV cache: {total_time:.6f} seconds\")"
      ],
      "metadata": {
        "id": "DbqNCmzTEDMX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\"This block implements inference with FullCache (baseline KV cache)\"\"\"\n",
        "import time\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "# Initialize the model\n",
        "model_config = WayformerTrainingConfig()\n",
        "model = WayformerCache(model_config, action_space)\n",
        "\n",
        "# Get a test scenario\n",
        "test_scenario = next(dataloader.simulator_state_generator(config=config))\n",
        "\n",
        "# Number of autoregressive steps for inference benchmarking\n",
        "num_autoregressive_steps = 1024  # Or any desired number of steps\n",
        "\n",
        "# Initialize the state with the test scenario\n",
        "state = test_scenario\n",
        "\n",
        "# Pad the state to have sufficient timesteps\n",
        "initial_timestep = state.timestep\n",
        "required_timesteps = initial_timestep + num_autoregressive_steps + 1  # +1 for zero-based indexing\n",
        "state = pad_state_to_num_timesteps(state, required_timesteps)\n",
        "\n",
        "# Initialize the history buffer for motion tokens\n",
        "history_motion_tokens = []\n",
        "\n",
        "total_time = 0.0\n",
        "\n",
        "# Initialize the FullCache for each decoder layer\n",
        "cache = {}\n",
        "for idx, decoder_layer in enumerate(model.decoder_layers):\n",
        "    cache_key = f'decoder_layer_{idx}_self_attn'\n",
        "    cache[cache_key] = FullCache(\n",
        "        max_cache_length=num_autoregressive_steps,\n",
        "        num_heads=model_config.tx_num_heads,\n",
        "        depth=model_config.hidden_size // model_config.tx_num_heads\n",
        "    )\n",
        "\n",
        "for step in range(num_autoregressive_steps):\n",
        "    # Process the current state to get model input\n",
        "    test_data = process_waymax_data(state, action_space)\n",
        "\n",
        "    # Update history buffers\n",
        "    if step == 0:\n",
        "        # For the first step, initialize with zeros or appropriate start tokens\n",
        "        initial_motion_token = 0  # Replace with your start token if applicable\n",
        "        history_motion_tokens.append(initial_motion_token)\n",
        "    else:\n",
        "        # Append the predicted motion token from the previous step\n",
        "        history_motion_tokens.append(predicted_motion_token)\n",
        "\n",
        "    # Prepare the model inputs, including the history\n",
        "    # With FullCache, we process one token at a time\n",
        "    model_input = prepare_model_input_with_history(\n",
        "        [test_data],\n",
        "        history_motion_tokens[-1:],  # Use only the last token\n",
        "        T=1\n",
        "    )\n",
        "\n",
        "    # Convert model inputs to tensors\n",
        "    model_inputs = {k: tf.convert_to_tensor(v, dtype=tf.float32) for k, v in model_input.items()}\n",
        "    # Perform inference\n",
        "    # Measure time before inference\n",
        "    start_time = time.time()\n",
        "\n",
        "    # Pass the cache to the model\n",
        "    predictions = model(model_inputs, cache=cache, training=False)\n",
        "\n",
        "    # Measure time after inference\n",
        "    end_time = time.time()\n",
        "    # Calculate elapsed time\n",
        "    elapsed_time = end_time - start_time\n",
        "    total_time += elapsed_time\n",
        "    # Print the timing result\n",
        "    print(f\"Auto-regressive inference time for step {step}: {elapsed_time:.6f} seconds\")\n",
        "\n",
        "    # Get the predicted motion token from logits\n",
        "    # 'logits' has shape (batch_size, 1, vocab_size)\n",
        "    # We need the predicted token for the current step\n",
        "    predicted_motion_token = tf.argmax(predictions['logits'], axis=-1).numpy()[0, 0]\n",
        "\n",
        "    # Convert motion token to acceleration\n",
        "    acceleration = action_space.continuous_acceleration(predicted_motion_token)  # Shape: (2,)\n",
        "\n",
        "    # Update the position and velocity using Verlet integration\n",
        "    ego_index = 0  # Assuming ego vehicle is at index 0\n",
        "\n",
        "    # Get current position and velocity from the state\n",
        "    current_position = np.array([\n",
        "        state.sim_trajectory.x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    current_velocity = np.array([\n",
        "        state.sim_trajectory.vel_x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.vel_y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    # Use the verlet_update function to compute the next position and velocity\n",
        "    new_position, new_velocity = action_space.verlet_update(current_position, current_velocity, acceleration)\n",
        "\n",
        "    # Optionally compute the new yaw angle based on the velocity vector\n",
        "    next_yaw = np.arctan2(new_velocity[1], new_velocity[0])  # Calculate yaw from velocity components\n",
        "\n",
        "    # Update the state with the predicted values\n",
        "    state = state.replace(\n",
        "        sim_trajectory=state.sim_trajectory.replace(\n",
        "            x=state.sim_trajectory.x.at[ego_index, state.timestep + 1].set(new_position[0]),\n",
        "            y=state.sim_trajectory.y.at[ego_index, state.timestep + 1].set(new_position[1]),\n",
        "            vel_x=state.sim_trajectory.vel_x.at[ego_index, state.timestep + 1].set(new_velocity[0]),\n",
        "            vel_y=state.sim_trajectory.vel_y.at[ego_index, state.timestep + 1].set(new_velocity[1]),\n",
        "            yaw=state.sim_trajectory.yaw.at[ego_index, state.timestep + 1].set(next_yaw),\n",
        "            valid=state.sim_trajectory.valid.at[ego_index, state.timestep + 1].set(True)\n",
        "        ),\n",
        "        timestep=state.timestep + 1  # Advance the timestep\n",
        "    )\n",
        "\n",
        "    # Break the loop if we've reached the end of available timesteps\n",
        "    if state.timestep >= state.sim_trajectory.x.shape[1] - 1:\n",
        "        break\n",
        "\n",
        "# Print total time\n",
        "print(f\"Total auto-regressive inference time with FullCache: {total_time:.6f} seconds\")\n",
        "# After the inference loop\n",
        "memory_usage_bytes = get_cache_memory_usage(cache)\n",
        "memory_usage_mb = memory_usage_bytes / (1024 ** 2)\n",
        "print(f\"Total KV cache memory usage: {memory_usage_mb:.2f} MB\")"
      ],
      "metadata": {
        "id": "m6UY7pRHQhcU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\"This block implements inference with SlidingWindowCache\"\"\"\n",
        "import time\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "# Initialize the model\n",
        "model_config = WayformerTrainingConfig()\n",
        "model = WayformerCache(model_config, action_space)\n",
        "\n",
        "# Get a test scenario\n",
        "test_scenario = next(dataloader.simulator_state_generator(config=config))\n",
        "\n",
        "# Number of autoregressive steps for inference benchmarking\n",
        "num_autoregressive_steps = 1024  # Or any desired number of steps\n",
        "\n",
        "# Initialize the state with the test scenario\n",
        "state = test_scenario\n",
        "\n",
        "# Pad the state to have sufficient timesteps\n",
        "initial_timestep = state.timestep\n",
        "required_timesteps = initial_timestep + num_autoregressive_steps + 1  # +1 for zero-based indexing\n",
        "state = pad_state_to_num_timesteps(state, required_timesteps)\n",
        "\n",
        "# Initialize the history buffer for motion tokens\n",
        "history_motion_tokens = []\n",
        "\n",
        "total_time = 0.0\n",
        "\n",
        "# Initialize the SlidingWindowCache for each decoder layer\n",
        "window_size = 10  # Set the desired window size\n",
        "cache = {}\n",
        "for idx, decoder_layer in enumerate(model.decoder_layers):\n",
        "    cache_key = f'decoder_layer_{idx}_self_attn'\n",
        "    cache[cache_key] = SlidingWindowCache(\n",
        "        window_size=window_size,\n",
        "        num_heads=model_config.tx_num_heads,\n",
        "        head_dim=model_config.hidden_size // model_config.tx_num_heads\n",
        "    )\n",
        "\n",
        "for step in range(num_autoregressive_steps):\n",
        "    # Process the current state to get model input\n",
        "    test_data = process_waymax_data(state, action_space)\n",
        "\n",
        "    # Update history buffers\n",
        "    if step == 0:\n",
        "        # For the first step, initialize with zeros or appropriate start tokens\n",
        "        initial_motion_token = 0  # Replace with your start token if applicable\n",
        "        history_motion_tokens.append(initial_motion_token)\n",
        "    else:\n",
        "        # Append the predicted motion token from the previous step\n",
        "        history_motion_tokens.append(predicted_motion_token)\n",
        "\n",
        "    # Prepare the model inputs, including the history\n",
        "    # With SlidingWindowCache, we only need the last token\n",
        "    model_input = prepare_model_input_with_history(\n",
        "        [test_data],\n",
        "        history_motion_tokens[-1:],  # Use only the last token\n",
        "        T=1\n",
        "    )\n",
        "\n",
        "    # Convert model inputs to tensors\n",
        "    model_inputs = {k: tf.convert_to_tensor(v, dtype=tf.float32) for k, v in model_input.items()}\n",
        "\n",
        "    # Perform inference\n",
        "    # Measure time before inference\n",
        "    start_time = time.time()\n",
        "\n",
        "    # Pass the cache to the model\n",
        "    predictions = model(model_inputs, cache=cache, training=False)\n",
        "\n",
        "    # Measure time after inference\n",
        "    end_time = time.time()\n",
        "    # Calculate elapsed time\n",
        "    elapsed_time = end_time - start_time\n",
        "    total_time += elapsed_time\n",
        "    # Print the timing result\n",
        "    print(f\"Auto-regressive inference time for step {step}: {elapsed_time:.6f} seconds\")\n",
        "\n",
        "    # Get the predicted motion token from logits\n",
        "    # 'logits' has shape (batch_size, 1, vocab_size)\n",
        "    # We need the predicted token for the current step\n",
        "    predicted_motion_token = tf.argmax(predictions['logits'], axis=-1).numpy()[0, 0]\n",
        "\n",
        "    # Convert motion token to acceleration\n",
        "    acceleration = action_space.continuous_acceleration(predicted_motion_token)  # Shape: (2,)\n",
        "\n",
        "    # Update the position and velocity using Verlet integration\n",
        "    ego_index = 0  # Assuming ego vehicle is at index 0\n",
        "\n",
        "    # Get current position and velocity from the state\n",
        "    current_position = np.array([\n",
        "        state.sim_trajectory.x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    current_velocity = np.array([\n",
        "        state.sim_trajectory.vel_x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.vel_y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    # Use the verlet_update function to compute the next position and velocity\n",
        "    new_position, new_velocity = action_space.verlet_update(current_position, current_velocity, acceleration)\n",
        "\n",
        "    # Optionally compute the new yaw angle based on the velocity vector\n",
        "    next_yaw = np.arctan2(new_velocity[1], new_velocity[0])  # Calculate yaw from velocity components\n",
        "\n",
        "    # Update the state with the predicted values\n",
        "    state = state.replace(\n",
        "        sim_trajectory=state.sim_trajectory.replace(\n",
        "            x=state.sim_trajectory.x.at[ego_index, state.timestep + 1].set(new_position[0]),\n",
        "            y=state.sim_trajectory.y.at[ego_index, state.timestep + 1].set(new_position[1]),\n",
        "            vel_x=state.sim_trajectory.vel_x.at[ego_index, state.timestep + 1].set(new_velocity[0]),\n",
        "            vel_y=state.sim_trajectory.vel_y.at[ego_index, state.timestep + 1].set(new_velocity[1]),\n",
        "            yaw=state.sim_trajectory.yaw.at[ego_index, state.timestep + 1].set(next_yaw),\n",
        "            valid=state.sim_trajectory.valid.at[ego_index, state.timestep + 1].set(True)\n",
        "        ),\n",
        "        timestep=state.timestep + 1  # Advance the timestep\n",
        "    )\n",
        "\n",
        "    # Break the loop if we've reached the end of available timesteps\n",
        "    if state.timestep >= state.sim_trajectory.x.shape[1] - 1:\n",
        "        break\n",
        "\n",
        "# Print total time\n",
        "print(f\"Total auto-regressive inference time with SlidingWindowCache: {total_time:.6f} seconds\")\n",
        "# After the inference loop\n",
        "memory_usage_bytes = get_cache_memory_usage(cache)\n",
        "memory_usage_mb = memory_usage_bytes / (1024 ** 2)\n",
        "print(f\"Total KV cache memory usage: {memory_usage_mb:.2f} MB\")"
      ],
      "metadata": {
        "id": "uGnnLzbLR4QY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\"This block implements inference with FullCache (baseline KV cache) with pre-filled cache using random data\"\"\"\n",
        "import time\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "# Initialize the model\n",
        "model_config = WayformerTrainingConfig()\n",
        "model = WayformerCache(model_config, action_space)\n",
        "\n",
        "# Get a test scenario\n",
        "test_scenario = next(dataloader.simulator_state_generator(config=config))\n",
        "\n",
        "# Define the maximum cache length\n",
        "max_cache_length = 16384  # Adjust as needed\n",
        "start_lenth = 16300\n",
        "\n",
        "# Define the number of steps to pre-fill the cache (e.g., half of max_cache_length)\n",
        "pre_fill_steps = start_lenth  # max_cache_length steps\n",
        "\n",
        "# Number of autoregressive steps for inference benchmarking\n",
        "num_autoregressive_steps = max_cache_length  # 16384 steps total\n",
        "\n",
        "# Initialize the state with the test scenario\n",
        "state = test_scenario\n",
        "\n",
        "# Pad the state to have sufficient timesteps\n",
        "initial_timestep = state.timestep\n",
        "required_timesteps = initial_timestep + num_autoregressive_steps + 1  # +1 for zero-based indexing\n",
        "state = pad_state_to_num_timesteps(state, required_timesteps)\n",
        "\n",
        "# Initialize the history buffer for motion tokens\n",
        "history_motion_tokens = []\n",
        "\n",
        "total_time = 0.0\n",
        "\n",
        "# Initialize the FullCache for each decoder layer\n",
        "cache = {}\n",
        "for idx, decoder_layer in enumerate(model.decoder_layers):\n",
        "    cache_key = f'decoder_layer_{idx}_self_attn'\n",
        "    cache[cache_key] = FullCache(\n",
        "        max_cache_length=max_cache_length,\n",
        "        num_heads=model_config.tx_num_heads,\n",
        "        depth=model_config.hidden_size // model_config.tx_num_heads\n",
        "    )\n",
        "\n",
        "# Function to pre-fill the cache with random data\n",
        "def pre_fill_cache_with_random(cache, pre_fill_steps):\n",
        "    \"\"\"\n",
        "    Pre-fills the cache with random key and value tensors.\n",
        "\n",
        "    Args:\n",
        "        cache (dict): Dictionary containing cache instances for each decoder layer.\n",
        "        pre_fill_steps (int): Number of tokens to pre-fill in the cache.\n",
        "    \"\"\"\n",
        "    print(f\"Starting pre-fill of the cache with {pre_fill_steps} steps using random data...\")\n",
        "    for cache_key, cache_instance in cache.items():\n",
        "        # Ensure cache is initialized\n",
        "        if cache_instance.key is None:\n",
        "            # Assume batch_size=1\n",
        "            batch_size = 1\n",
        "            random_key = tf.random.normal(\n",
        "                [batch_size, cache_instance.num_heads, pre_fill_steps, cache_instance.depth],\n",
        "                dtype=tf.float32\n",
        "            )\n",
        "            random_value = tf.random.normal(\n",
        "                [batch_size, cache_instance.num_heads, pre_fill_steps, cache_instance.depth],\n",
        "                dtype=tf.float32\n",
        "            )\n",
        "\n",
        "            # Initialize the cache's key and value with zeros first\n",
        "            cache_instance.key = tf.Variable(\n",
        "                tf.zeros([batch_size, cache_instance.num_heads, cache_instance.max_cache_length, cache_instance.depth], dtype=tf.float32),\n",
        "                trainable=False\n",
        "            )\n",
        "            cache_instance.value = tf.Variable(\n",
        "                tf.zeros([batch_size, cache_instance.num_heads, cache_instance.max_cache_length, cache_instance.depth], dtype=tf.float32),\n",
        "                trainable=False\n",
        "            )\n",
        "\n",
        "            # Assign the random data to the first pre_fill_steps positions\n",
        "            cache_instance.key[:, :, :pre_fill_steps, :].assign(random_key)\n",
        "            cache_instance.value[:, :, :pre_fill_steps, :].assign(random_value)\n",
        "\n",
        "            # Update current_length\n",
        "            cache_instance.current_length = pre_fill_steps\n",
        "        else:\n",
        "            # If cache is already initialized, append random data\n",
        "            for _ in range(pre_fill_steps):\n",
        "                random_key = tf.random.normal(\n",
        "                    [1, cache_instance.num_heads, 1, cache_instance.depth],\n",
        "                    dtype=tf.float32\n",
        "                )\n",
        "                random_value = tf.random.normal(\n",
        "                    [1, cache_instance.num_heads, 1, cache_instance.depth],\n",
        "                    dtype=tf.float32\n",
        "                )\n",
        "                cache_instance.update(random_key, random_value)\n",
        "\n",
        "    print(\"Pre-fill completed.\\n\")\n",
        "\n",
        "# Function to get cache memory usage (for illustration purposes)\n",
        "def get_cache_memory_usage(cache):\n",
        "    \"\"\"\n",
        "    Estimates the memory usage of the cache.\n",
        "\n",
        "    Args:\n",
        "        cache (dict): Dictionary containing cache instances for each decoder layer.\n",
        "\n",
        "    Returns:\n",
        "        int: Total memory usage in bytes.\n",
        "    \"\"\"\n",
        "    total_bytes = 0\n",
        "    for cache_instance in cache.values():\n",
        "        if cache_instance.key is not None:\n",
        "            total_bytes += cache_instance.key.numpy().nbytes\n",
        "        if cache_instance.value is not None:\n",
        "            total_bytes += cache_instance.value.numpy().nbytes\n",
        "    return total_bytes\n",
        "\n",
        "# Pre-fill the cache with random data\n",
        "pre_fill_cache_with_random(cache, pre_fill_steps)\n",
        "\n",
        "# Define the number of main inference steps\n",
        "main_steps = num_autoregressive_steps - pre_fill_steps  # 16384 - 8192 = 8192 steps\n",
        "\n",
        "print(f\"Starting main inference loop with {main_steps} steps...\")\n",
        "\n",
        "for step in range(main_steps):\n",
        "    # Process the current state to get model input\n",
        "    test_data = process_waymax_data(state, action_space)\n",
        "\n",
        "    # Append the predicted motion token from the previous step\n",
        "    # Note: For pre-fill, history_motion_tokens already contains pre_fill_steps tokens\n",
        "    history_motion_tokens.append(predicted_motion_token)\n",
        "\n",
        "    # Prepare the model inputs, including the latest token\n",
        "    # With FullCache, we process one token at a time\n",
        "    model_input = prepare_model_input_with_history(\n",
        "        [test_data],\n",
        "        history_motion_tokens[-1:],  # Use only the last token\n",
        "        T=1\n",
        "    )\n",
        "\n",
        "    # Convert model inputs to tensors\n",
        "    model_inputs = {k: tf.convert_to_tensor(v, dtype=tf.float32) for k, v in model_input.items()}\n",
        "    # Print input shape\n",
        "    print(f\"Model input shape for step {step}: {model_inputs['motion_tokens'].shape}\")\n",
        "\n",
        "    # Perform inference\n",
        "    # Measure time before inference\n",
        "    start_time = time.time()\n",
        "\n",
        "    # Pass the cache to the model\n",
        "    predictions = model(model_inputs, cache=cache, training=False)\n",
        "\n",
        "    # Measure time after inference\n",
        "    end_time = time.time()\n",
        "    # Calculate elapsed time\n",
        "    elapsed_time = end_time - start_time\n",
        "    total_time += elapsed_time\n",
        "    # Print the timing result\n",
        "    print(f\"Auto-regressive inference time for step {step}: {elapsed_time:.6f} seconds\")\n",
        "\n",
        "    # Get the predicted motion token from logits\n",
        "    # 'logits' has shape (batch_size, 1, vocab_size)\n",
        "    # We need the predicted token for the current step\n",
        "    predicted_motion_token = tf.argmax(predictions['logits'], axis=-1).numpy()[0, 0]\n",
        "\n",
        "    # Convert motion token to acceleration\n",
        "    acceleration = action_space.continuous_acceleration(predicted_motion_token)  # Shape: (2,)\n",
        "\n",
        "    # Update the position and velocity using Verlet integration\n",
        "    ego_index = 0  # Assuming ego vehicle is at index 0\n",
        "\n",
        "    # Get current position and velocity from the state\n",
        "    current_position = np.array([\n",
        "        state.sim_trajectory.x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    current_velocity = np.array([\n",
        "        state.sim_trajectory.vel_x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.vel_y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    # Use the verlet_update function to compute the next position and velocity\n",
        "    new_position, new_velocity = action_space.verlet_update(current_position, current_velocity, acceleration)\n",
        "\n",
        "    # Optionally compute the new yaw angle based on the velocity vector\n",
        "    next_yaw = np.arctan2(new_velocity[1], new_velocity[0])  # Calculate yaw from velocity components\n",
        "\n",
        "    # Update the state with the predicted values\n",
        "    state = state.replace(\n",
        "        sim_trajectory=state.sim_trajectory.replace(\n",
        "            x=state.sim_trajectory.x.at[ego_index, state.timestep + 1].set(new_position[0]),\n",
        "            y=state.sim_trajectory.y.at[ego_index, state.timestep + 1].set(new_position[1]),\n",
        "            vel_x=state.sim_trajectory.vel_x.at[ego_index, state.timestep + 1].set(new_velocity[0]),\n",
        "            vel_y=state.sim_trajectory.vel_y.at[ego_index, state.timestep + 1].set(new_velocity[1]),\n",
        "            yaw=state.sim_trajectory.yaw.at[ego_index, state.timestep + 1].set(next_yaw),\n",
        "            valid=state.sim_trajectory.valid.at[ego_index, state.timestep + 1].set(True)\n",
        "        ),\n",
        "        timestep=state.timestep + 1  # Advance the timestep\n",
        "    )\n",
        "\n",
        "    # Optionally, print progress\n",
        "    if (step + 1) % 1000 == 0:\n",
        "        print(f\"Completed {step + 1}/{main_steps} inference steps.\\n\")\n",
        "\n",
        "print(\"Main inference loop completed.\\n\")\n",
        "\n",
        "# Print total time\n",
        "print(f\"Total auto-regressive inference time with FullCache: {total_time:.6f} seconds\")\n",
        "# After the inference loop\n",
        "memory_usage_bytes = get_cache_memory_usage(cache)\n",
        "memory_usage_mb = memory_usage_bytes / (1024 ** 2)\n",
        "print(f\"Total KV cache memory usage: {memory_usage_mb:.2f} MB\")"
      ],
      "metadata": {
        "id": "qdhNZzPaR6dG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\"This block implements inference with SlidingWindowCache with pre-filled cache using random data\"\"\"\n",
        "import time\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "# Initialize the model\n",
        "model_config = WayformerTrainingConfig()\n",
        "action_space = VerletActionSpace()  # Ensure this is defined appropriately\n",
        "model = WayformerCache(model_config, action_space)\n",
        "\n",
        "# Get a test scenario\n",
        "test_scenario = next(dataloader.simulator_state_generator(config=config))  # Ensure dataloader and config are defined\n",
        "\n",
        "# Define the maximum cache length\n",
        "window_size = 80  # Adjust based on your performance and memory requirements\n",
        "\n",
        "# Define the number of steps to pre-fill the cache (exactly window_size for SlidingWindowCache)\n",
        "pre_fill_steps = window_size  # 4096 steps\n",
        "\n",
        "# Number of autoregressive steps for inference benchmarking\n",
        "num_autoregressive_steps = 16384  # Total steps including pre-fill\n",
        "start_step = 16300\n",
        "\n",
        "# Initialize the state with the test scenario\n",
        "state = test_scenario\n",
        "\n",
        "# Pad the state to have sufficient timesteps\n",
        "initial_timestep = state.timestep\n",
        "required_timesteps = initial_timestep + num_autoregressive_steps + 1  # +1 for zero-based indexing\n",
        "state = pad_state_to_num_timesteps(state, required_timesteps)  # Ensure this function is defined\n",
        "\n",
        "# Initialize the history buffer for motion tokens\n",
        "history_motion_tokens = []\n",
        "\n",
        "total_time = 0.0\n",
        "\n",
        "# Initialize the SlidingWindowCache for each decoder layer\n",
        "cache = {}\n",
        "for idx, decoder_layer in enumerate(model.decoder_layers):\n",
        "    cache_key = f'decoder_layer_{idx}_self_attn'\n",
        "    cache[cache_key] = SlidingWindowCache(\n",
        "        window_size=window_size,\n",
        "        num_heads=model_config.tx_num_heads,\n",
        "        head_dim=model_config.hidden_size // model_config.tx_num_heads\n",
        "    )\n",
        "\n",
        "# Function to pre-fill the SlidingWindowCache with random data\n",
        "def pre_fill_cache_with_random(cache, pre_fill_steps):\n",
        "    \"\"\"\n",
        "    Pre-fills the SlidingWindowCache with random key and value tensors.\n",
        "\n",
        "    Args:\n",
        "        cache (dict): Dictionary containing cache instances for each decoder layer.\n",
        "        pre_fill_steps (int): Number of tokens to pre-fill in the cache.\n",
        "    \"\"\"\n",
        "    print(f\"Starting pre-fill of the cache with {pre_fill_steps} tokens using random data...\")\n",
        "    for cache_key, cache_instance in cache.items():\n",
        "        # Assume batch_size=1\n",
        "        batch_size = 1\n",
        "\n",
        "        # Generate random key and value tensors with shape [batch_size, num_heads, 1, head_dim]\n",
        "        random_keys = tf.random.normal(\n",
        "            [batch_size, cache_instance.num_heads, pre_fill_steps, cache_instance.head_dim],\n",
        "            dtype=tf.float32\n",
        "        )\n",
        "        random_values = tf.random.normal(\n",
        "            [batch_size, cache_instance.num_heads, pre_fill_steps, cache_instance.head_dim],\n",
        "            dtype=tf.float32\n",
        "        )\n",
        "\n",
        "        # Initialize the cache's key and value if not already initialized\n",
        "        if cache_instance.key_buffer is None:\n",
        "            # Initialize with zeros\n",
        "            cache_instance.key_buffer = tf.Variable(\n",
        "                tf.zeros([batch_size, cache_instance.num_heads, cache_instance.window_size, cache_instance.head_dim], dtype=tf.float32),\n",
        "                trainable=False\n",
        "            )\n",
        "            cache_instance.value_buffer = tf.Variable(\n",
        "                tf.zeros([batch_size, cache_instance.num_heads, cache_instance.window_size, cache_instance.head_dim], dtype=tf.float32),\n",
        "                trainable=False\n",
        "            )\n",
        "            cache_instance.current_size.assign(0)\n",
        "            cache_instance.start_idx.assign(0)\n",
        "            cache_instance.end_idx.assign(0)\n",
        "\n",
        "        # Assign random keys and values to the cache buffer\n",
        "        for i in range(pre_fill_steps):\n",
        "            key = random_keys[:, :, i, :]  # Shape: [1, num_heads, head_dim]\n",
        "            value = random_values[:, :, i, :]  # Shape: [1, num_heads, head_dim]\n",
        "            # Expand dimensions to match [batch_size, num_heads, seq_len_k, head_dim]\n",
        "            key = tf.expand_dims(key, axis=2)  # Now [1, num_heads, 1, head_dim]\n",
        "            value = tf.expand_dims(value, axis=2)  # Now [1, num_heads, 1, head_dim]\n",
        "            cache_instance.update(key, value)\n",
        "    print(\"Pre-fill with random data completed.\\n\")\n",
        "\n",
        "# Function to get cache memory usage (for illustration purposes)\n",
        "def get_cache_memory_usage(cache):\n",
        "    \"\"\"\n",
        "    Estimates the memory usage of the cache.\n",
        "\n",
        "    Args:\n",
        "        cache (dict): Dictionary containing cache instances for each decoder layer.\n",
        "\n",
        "    Returns:\n",
        "        int: Total memory usage in bytes.\n",
        "    \"\"\"\n",
        "    total_bytes = 0\n",
        "    for cache_instance in cache.values():\n",
        "        if cache_instance.key_buffer is not None:\n",
        "            total_bytes += cache_instance.key_buffer.numpy().nbytes\n",
        "        if cache_instance.value_buffer is not None:\n",
        "            total_bytes += cache_instance.value_buffer.numpy().nbytes\n",
        "    return total_bytes\n",
        "\n",
        "# Pre-fill the cache with random data\n",
        "pre_fill_cache_with_random(cache, pre_fill_steps)\n",
        "\n",
        "# Define the number of main inference steps\n",
        "main_steps = num_autoregressive_steps - start_step  # e.g., 16384 - 16000 = 384\n",
        "\n",
        "print(f\"Starting main inference loop with {main_steps} steps...\")\n",
        "\n",
        "for step in range(main_steps):\n",
        "    # Process the current state to get model input\n",
        "    test_data = process_waymax_data(state, action_space)  # Ensure this function is defined\n",
        "\n",
        "    # Append the predicted motion token from the previous step\n",
        "    # For the first step after pre-fill, you might need an initial token\n",
        "    if step == 0:\n",
        "        initial_motion_token = 0  # Replace with your start token if applicable\n",
        "        history_motion_tokens.append(initial_motion_token)\n",
        "    else:\n",
        "        history_motion_tokens.append(predicted_motion_token)\n",
        "\n",
        "    # Prepare the model inputs, including the latest token\n",
        "    # With SlidingWindowCache, we process one token at a time\n",
        "    model_input = prepare_model_input_with_history(\n",
        "        [test_data],\n",
        "        history_motion_tokens[-1:],  # Use only the last token\n",
        "        T=1  # Since we're doing one step at a time\n",
        "    )\n",
        "\n",
        "    # Convert model inputs to tensors\n",
        "    model_inputs = {k: tf.convert_to_tensor(v, dtype=tf.float32) for k, v in model_input.items()}\n",
        "\n",
        "    # Optionally, print input shape\n",
        "    print(f\"Model input shape for step {step}: {model_inputs['motion_tokens'].shape}\")\n",
        "\n",
        "    # Perform inference\n",
        "    # Measure time before inference\n",
        "    start_time = time.time()\n",
        "\n",
        "    # Pass the cache to the model\n",
        "    predictions = model(model_inputs, cache=cache, training=False)\n",
        "\n",
        "    # Measure time after inference\n",
        "    end_time = time.time()\n",
        "    # Calculate elapsed time\n",
        "    elapsed_time = end_time - start_time\n",
        "    total_time += elapsed_time\n",
        "    # Print the timing result\n",
        "    print(f\"Auto-regressive inference time for step {step}: {elapsed_time:.6f} seconds\")\n",
        "\n",
        "    # Get the predicted motion token from logits\n",
        "    # 'logits' has shape (batch_size, 1, vocab_size)\n",
        "    # We need the predicted token for the current step\n",
        "    predicted_motion_token = tf.argmax(predictions['logits'], axis=-1).numpy()[0, 0]\n",
        "\n",
        "    # Convert motion token to acceleration\n",
        "    acceleration = action_space.continuous_acceleration(predicted_motion_token)  # Shape: (2,)\n",
        "\n",
        "    # Update the position and velocity using Verlet integration\n",
        "    ego_index = 0  # Assuming ego vehicle is at index 0\n",
        "\n",
        "    # Get current position and velocity from the state\n",
        "    current_position = np.array([\n",
        "        state.sim_trajectory.x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    current_velocity = np.array([\n",
        "        state.sim_trajectory.vel_x[ego_index, state.timestep],\n",
        "        state.sim_trajectory.vel_y[ego_index, state.timestep]\n",
        "    ])\n",
        "\n",
        "    # Use the verlet_update function to compute the next position and velocity\n",
        "    new_position, new_velocity = action_space.verlet_update(current_position, current_velocity, acceleration)\n",
        "\n",
        "    # Optionally compute the new yaw angle based on the velocity vector\n",
        "    next_yaw = np.arctan2(new_velocity[1], new_velocity[0])  # Calculate yaw from velocity components\n",
        "\n",
        "    # Update the state with the predicted values\n",
        "    state = state.replace(\n",
        "        sim_trajectory=state.sim_trajectory.replace(\n",
        "            x=state.sim_trajectory.x.at[ego_index, state.timestep + 1].set(new_position[0]),\n",
        "            y=state.sim_trajectory.y.at[ego_index, state.timestep + 1].set(new_position[1]),\n",
        "            vel_x=state.sim_trajectory.vel_x.at[ego_index, state.timestep + 1].set(new_velocity[0]),\n",
        "            vel_y=state.sim_trajectory.vel_y.at[ego_index, state.timestep + 1].set(new_velocity[1]),\n",
        "            yaw=state.sim_trajectory.yaw.at[ego_index, state.timestep + 1].set(next_yaw),\n",
        "            valid=state.sim_trajectory.valid.at[ego_index, state.timestep + 1].set(True)\n",
        "        ),\n",
        "        timestep=state.timestep + 1  # Advance the timestep\n",
        "    )\n",
        "\n",
        "    # Optionally, print progress every 1000 steps\n",
        "    if (step + 1) % 1000 == 0:\n",
        "        print(f\"Completed {step + 1}/{main_steps} inference steps.\\n\")\n",
        "\n",
        "print(\"Main inference loop completed.\\n\")\n",
        "\n",
        "# Print total time\n",
        "print(f\"Total auto-regressive inference time with SlidingWindowCache: {total_time:.6f} seconds\")\n",
        "# After the inference loop\n",
        "memory_usage_bytes = get_cache_memory_usage(cache)\n",
        "memory_usage_mb = memory_usage_bytes / (1024 ** 2)\n",
        "print(f\"Total KV cache memory usage: {memory_usage_mb:.2f} MB\")"
      ],
      "metadata": {
        "id": "aY8VumLuJSnC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "l3KYVBXmMRud"
      },
      "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
}