Fix broken plots

mixtape/core/static/mixtape/core/insights/clusteringTimeline.js

@@ -59,7 +59,6 @@ document.addEventListener('alpine:init', () => {
           const step = this.$store.insights.currentStep;
           const xLine = [step, step];
           const yLine = [yMin - 0.5, yMax + 0.5];
-          // Use restyle: arguments must be arrays-of-arrays per selected trace
           Plotly.restyle(this.$el, { x: [xLine], y: [yLine] }, [1]);
         }
       });
@@ -99,7 +98,7 @@ document.addEventListener('alpine:init', () => {
         : 1;
       this.data = [
         {
-          agentStepClusters: agentStepClusters,
+          z: agentStepClusters,
           type: 'heatmap',
           showscale: false,
           colorscale: 'YlGnBu',

mixtape/core/static/mixtape/core/insights/insightsStore.js

@@ -21,6 +21,43 @@ document.addEventListener('alpine:init', () => {
     // Runtime values
     currentStep: 0,
 
+    // Get active timeline key steps based on global/episode mode
+    getActiveTimelineKeySteps() {
+      const settings = Alpine.store('settings');
+      if (!settings) return [];
+      const useGlobal = settings.useGlobalTimeline;
+      const epIdx = settings.timelineEpisodeIdx;
+
+      if (useGlobal && Array.isArray(this.timelineKeySteps)) {
+        // Global mode: merge all episodes' key steps, sort by reward, take top 40
+        const allSteps = [];
+        this.timelineKeySteps.forEach((episodeSteps, episodeIdx) => {
+          if (Array.isArray(episodeSteps)) {
+            episodeSteps.forEach(step => {
+              allSteps.push({
+                ...step,
+                episodeIdx: episodeIdx,
+                episodeId: this.episodeIds?.[episodeIdx] || episodeIdx,
+              });
+            });
+          }
+        });
+        // Sort by total_rewards descending, take top 40, then sort by step number
+        const topSteps = allSteps
+          .sort((a, b) => b.total_rewards - a.total_rewards)
+          .slice(0, 40)
+          .sort((a, b) => {
+            // Sort by episode first, then by step number within episode
+            if (a.episodeIdx !== b.episodeIdx) return a.episodeIdx - b.episodeIdx;
+            return a.number - b.number;
+          });
+        return topSteps;
+      } else {
+        // Episode mode: return selected episode's key steps
+        return this.timelineKeySteps?.[epIdx] || [];
+      }
+    },
+
     init() {
       // Parse JSON data
       this.parsedData = JSON.parse(

mixtape/core/static/mixtape/core/insights/replay.js

@@ -1,7 +1,14 @@
 document.addEventListener('alpine:init', () => {
   Alpine.data('replay', () => ({
     playing: false,
+    stepInput: 1,
     maxSteps: Alpine.store('insights').parsedData.max_steps - 1,
+    init() {
+      this.stepInput = this.$store.insights.currentStep + 1;
+      this.$watch('$store.insights.currentStep', (val) => {
+        this.stepInput = val + 1;
+      });
+    },
     stepForward() {
       if (!this.playing) {
         return;
@@ -14,8 +21,12 @@ document.addEventListener('alpine:init', () => {
         }, delay);
       } else {
         this.playing = false;
-        this.$store.insights.currentStep = 0;
       }
     },
+    jumpToStep() {
+      const val = Math.max(1, Math.min(this.maxSteps + 1, this.stepInput || 1));
+      this.stepInput = val;
+      this.$store.insights.currentStep = val - 1;
+    },
   }));
 });

mixtape/core/templates/core/insights/episode_info.html

@@ -1,41 +1,65 @@
-<section class="border border-base-300 rounded-lg p-3 h-full">
-  <div class="overflow-x-auto">
-    <table class="w-full text-sm">
-      <thead>
-        <tr>
-          <th class="text-left py-2 px-2 align-bottom" rowspan="2">Agents</th>
-          <template x-for="(ep, epIdx) in episodeSummaries" :key="epIdx">
-            <th scope="colgroup" class="text-center py-2 px-2 align-bottom font-medium" :colspan="2" x-text="`Episode ${ep.id}`"></th>
-          </template>
-        </tr>
-        <tr>
-          <template x-for="pair in episodePairs" :key="`h2-${pair.epIdx}-${pair.kind}`">
-            <th class="py-1 px-2 min-w-[6rem]"
-                :class="[
-                  pair.kind==='action' ? 'text-right' : 'text-left',
-                  (pair.kind==='action' && pair.epIdx>0) ? 'border-l' : ''
-                ].join(' ')"
-                x-text="pair.kind==='action' ? 'Action' : 'Reward'">
-            </th>
-          </template>
-        </tr>
-      </thead>
-      <tbody>
-        <template x-for="agent in $store.insights.uniqueAgents" :key="agent">
-          <tr class="border-t">
-            <td class="py-2 px-2 font-medium" x-text="agent"></td>
-            <template x-for="pair in episodePairs" :key="`r-${agent}-${pair.epIdx}-${pair.kind}`">
-              <td class="py-2 px-2 whitespace-nowrap min-w-[6rem] tabular-nums"
-                  :class="[
-                    pair.kind==='action' ? 'text-right' : 'text-left',
-                    (pair.kind==='action' && pair.epIdx>0) ? 'border-l' : ''
-                  ].join(' ')"
-                  x-text="pair.kind==='action' ? $data.getAgentAction(pair.epIdx, agent) : $data.getAgentRewardStr(pair.epIdx, agent)">
-              </td>
+<section class="border border-base-300 rounded-lg p-3 h-full flex flex-col gap-3"
+         x-data="{
+           selectedEpisodes: [],
+           get allSelected() { return this.selectedEpisodes.length === (episodeSummaries?.length || 0); },
+           toggleAll() {
+             if (this.allSelected) {
+               this.selectedEpisodes = [];
+             } else {
+               this.selectedEpisodes = episodeSummaries?.map((_, i) => i) || [];
+             }
+           },
+           init() { this.selectedEpisodes = episodeSummaries?.map((_, i) => i) || []; }
+         }">
+  <div class="flex items-center justify-between flex-wrap gap-2">
+    <h3 class="text-sm font-medium">Step Details</h3>
+    <div x-show="episodeSummaries?.length > 1" x-cloak class="flex items-center gap-3 text-sm">
+      <label class="flex items-center gap-1 cursor-pointer">
+        <input type="checkbox" class="checkbox checkbox-sm"
+               :checked="allSelected"
+               @change="toggleAll()" />
+        <span>All</span>
+      </label>
+      <template x-for="(ep, idx) in episodeSummaries" :key="idx">
+        <label class="flex items-center gap-1 cursor-pointer">
+          <input type="checkbox" class="checkbox checkbox-sm"
+                 :value="idx"
+                 x-model.number="selectedEpisodes" />
+          <span x-text="`Ep ${ep.id}`"></span>
+        </label>
+      </template>
+    </div>
+  </div>
+  <div class="flex flex-wrap gap-4 overflow-y-auto">
+    <template x-for="(ep, epIdx) in episodeSummaries" :key="epIdx">
+      <div class="flex-1 min-w-[200px]" x-show="selectedEpisodes.includes(epIdx)">
+        <table class="text-sm w-full">
+          <thead>
+            <tr>
+              <th scope="colgroup" class="text-center py-2 px-2 font-medium border-b" colspan="3">
+                <span x-text="`Episode ${ep.id}`"></span>
+                <span x-show="$store.insights.currentStep >= (ep.steps || 0)" x-cloak
+                      class="ml-1 badge badge-xs badge-success"
+                      title="Episode completed">done</span>
+              </th>
+            </tr>
+            <tr class="border-b">
+              <th class="text-left py-1 px-2">Agent</th>
+              <th class="text-right py-1 px-2 min-w-[5rem]">Action</th>
+              <th class="text-left py-1 px-2 min-w-[5rem]">Reward</th>
+            </tr>
+          </thead>
+          <tbody>
+            <template x-for="agent in $store.insights.uniqueAgents" :key="`${epIdx}-${agent}`">
+              <tr class="border-t">
+                <td class="py-2 px-2 font-medium" x-text="agent"></td>
+                <td class="py-2 px-2 text-right whitespace-nowrap tabular-nums" x-text="$data.getAgentAction(epIdx, agent)"></td>
+                <td class="py-2 px-2 text-left whitespace-nowrap tabular-nums" x-text="$data.getAgentRewardStr(epIdx, agent)"></td>
+              </tr>
             </template>
-          </tr>
-        </template>
-      </tbody>
-    </table>
+          </tbody>
+        </table>
+      </div>
+    </template>
   </div>
 </section>

mixtape/core/templates/core/insights/media_viewer.html

@@ -1,4 +1,4 @@
-<section class="border border-base-300 rounded-lg p-3 flex flex-col gap-3">
+<section class="border border-base-300 rounded-lg p-3 flex flex-col gap-3 h-full">
   <div class="flex items-center justify-between">
     <h3 class="text-sm font-medium">Media Viewer</h3>
     <div class="flex items-center gap-2">
@@ -10,7 +10,7 @@ <h3 class="text-sm font-medium">Media Viewer</h3>
     </div>
   </div>
   <template x-if="$data.getEpisodeImageUrl(mediaEpisodeIdx)">
-    <img :src="$data.getEpisodeImageUrl(mediaEpisodeIdx)" alt="Episode image" />
+    <img class="m-auto" :src="$data.getEpisodeImageUrl(mediaEpisodeIdx)" alt="Episode image" />
   </template>
   <template x-if="!$data.getEpisodeImageUrl(mediaEpisodeIdx)">
     <div class="badge badge-dash badge-warning">No image available</div>

mixtape/core/templates/core/insights/partials/_tabs_overview.html

@@ -58,11 +58,19 @@
         <i class="ri-checkbox-multiple-blank-line"></i>
         <span class="font-semibold" x-text="$store.insights.episodeIds.length"></span>
       </span>
-      <span class="badge badge-soft badge-primary" :title="`${$store.insights.uniqueAgents.length} Agents`">
+      <span class="badge badge-soft badge-primary"
+            :title="`Agents:\n${$store.insights.uniqueAgents.join('\n')}`">
         <i class="ri-robot-2-line"></i>
         <span class="font-semibold" x-text="$store.insights.uniqueAgents.length"></span>
       </span>
-      <span class="badge badge-soft badge-primary" :title="`Avg Reward: ${avgRewardDisplay ?? '—'}`">
+      <span class="badge badge-soft badge-primary"
+            :title="(() => {
+              const rewards = $store.insights.parsedData?.episode_total_rewards || [];
+              const ids = $store.insights.episodeIds || [];
+              if (!rewards.length) return `Avg Reward: ${avgRewardDisplay ?? '—'}`;
+              const lines = rewards.map((r, i) => `Episode ${ids[i] ?? i}: ${typeof r === 'number' ? r.toFixed(2) : '—'}`);
+              return `Avg: ${avgRewardDisplay ?? '—'}\n${lines.join('\n')}`;
+            })()">
         <i class="ri-trophy-line"></i>
         <span class="font-semibold" x-text="avgRewardDisplay ?? '—'"></span>
       </span>
@@ -224,11 +232,16 @@ <h3 class="text-sm font-medium mb-2">Actions</h3>
     <div x-show="$store.settings.showTimeline" x-cloak>
       {% include "core/insights/timeline.html" %}
     </div>
-    <div class="flex flex-row gap-2">
-      <div x-show="$store.settings.showEpisodeInfo" x-cloak class="flex-1">
+    <div class="grid gap-2"
+         :class="{
+           'grid-cols-2': $store.settings.showEpisodeInfo && $store.settings.showMediaViewer,
+           'grid-cols-1': !$store.settings.showEpisodeInfo || !$store.settings.showMediaViewer
+         }">
+      <div x-show="$store.settings.showEpisodeInfo" x-cloak class="overflow-y-auto max-h-full"
+           :class="{ 'row-span-1': $store.settings.showMediaViewer }">
         {% include "core/insights/episode_info.html" %}
       </div>
-      <div x-show="$store.settings.showMediaViewer" x-cloak class="flex-2">
+      <div x-show="$store.settings.showMediaViewer" x-cloak>
         {% include "core/insights/media_viewer.html" %}
       </div>
     </div>

mixtape/core/templates/core/insights/replay_controls.html

@@ -3,19 +3,22 @@
     class="btn btn-ghost btn-sm"
     :disabled="$store.insights.currentStep === 0"
     @click="$store.insights.currentStep = 0"
+    title="Go to first step"
   >
     <i class="ri-rewind-start-line"></i>
   </button>
   <button
     class="btn btn-ghost btn-sm"
     :disabled="$store.insights.currentStep === 0"
     @click="$store.insights.currentStep--"
+    title="Previous step"
   >
     <i class="ri-rewind-line"></i>
   </button>
   <button
     class="btn btn-ghost btn-sm"
     @click="playing = !playing; if (playing) stepForward()"
+    :title="playing ? 'Pause' : 'Play'"
   >
     <i x-show="!playing" x-cloak class="ri-play-line"></i>
     <i x-show="playing" x-cloak class="ri-pause-line"></i>
@@ -24,13 +27,15 @@
     class="btn btn-ghost btn-sm"
     :disabled="$store.insights.currentStep === maxSteps"
     @click="$store.insights.currentStep++"
+    title="Next step"
   >
     <i class="ri-speed-line"></i>
   </button>
   <button
     class="btn btn-ghost btn-sm"
     :disabled="$store.insights.currentStep === maxSteps"
     @click="$store.insights.currentStep = maxSteps"
+    title="Go to last step"
   >
     <i class="ri-forward-end-line"></i>
   </button>
@@ -39,10 +44,12 @@
       x-cloak
       type="number"
       class="input input-sm validator w-16 text-center"
-      min="0"
-      :max="maxSteps"
-      x-model.number="$store.insights.currentStep"
+      min="1"
+      :max="maxSteps + 1"
+      x-model.number="stepInput"
+      @keydown.enter="jumpToStep()"
+      @blur="jumpToStep()"
     />
-    <span class="text-xs">/ <span x-text="maxSteps"></span></span>
+    <span class="text-xs">/ <span x-text="maxSteps + 1"></span></span>
   </div>
 </section>

mixtape/core/templates/core/insights/timeline.html

@@ -1,17 +1,21 @@
 <section x-cloak class="border border-base-300 rounded-lg flex">
   <ol class="items-center flex p-1 w-full">
     <li class="relative flex items-center"><div class="flex w-1 bg-primary h-6 ml-2"></div></li>
-    <template x-for="ts in $store.insights.timelineKeySteps[$store.settings.timelineEpisodeIdx]">
-      <li class="relative w-full flex even:flex-col odd:flex-col-reverse even:mb-5 odd:mt-5" :key="ts.number">
+    <template x-for="ts in $store.insights.getActiveTimelineKeySteps()">
+      <li class="relative w-full flex even:flex-col odd:flex-col-reverse even:mb-5 odd:mt-5" :key="ts.episodeIdx !== undefined ? `${ts.episodeIdx}-${ts.number}` : ts.number">
         <p class="text-sm text-center" x-text="ts.total_rewards.toFixed(2)"></p>
         <div class="flex items-center">
             <div class="flex w-full bg-primary h-0.5"></div>
             <button
               class="cursor-pointer z-10 flex items-center justify-center w-2 h-2 rounded-full ring-4 shrink-0"
               :class="$store.insights.currentStep === ts.number ? 'ring-[var(--color-primary-content)] bg-primary' : 'ring-[var(--color-primary)]'"
               @click="$store.insights.currentStep = ts.number"
-              :title="`Step ${ts.number} — Reward: ${ts.total_rewards.toFixed(2)}`"
-              :aria-label="`Step ${ts.number} Reward ${ts.total_rewards.toFixed(2)}`"
+              :title="ts.episodeId !== undefined
+                ? `Episode ${ts.episodeId} — Step ${ts.number} — Reward: ${ts.total_rewards.toFixed(2)}`
+                : `Step ${ts.number} — Reward: ${ts.total_rewards.toFixed(2)}`"
+              :aria-label="ts.episodeId !== undefined
+                ? `Episode ${ts.episodeId} Step ${ts.number} Reward ${ts.total_rewards.toFixed(2)}`
+                : `Step ${ts.number} Reward ${ts.total_rewards.toFixed(2)}`"
             >
             </button>
             <div class="flex w-full bg-primary h-0.5"></div>

mixtape/core/templates/core/shared/partials/_side_panel_overview.html

@@ -114,7 +114,7 @@ <h4 class="text-xs opacity-60 mb-1">Global</h4>
               <input type="checkbox" x-model.boolean="$store.settings.showEpisodeInfo"
                      class="checkbox checkbox-sm"
               />
-              Episode Info
+              Step Details
             </label>
             <label class="flex items-center gap-2 text-sm">
               <input type="checkbox" x-model.boolean="$store.settings.showMediaViewer"

mixtape/core/views.py

@@ -302,9 +302,15 @@ def insights(request: HttpRequest) -> HttpResponse:
     for episode_pk in episode_pks:
         insight_results = _episode_insights(episode_pk, group_by_episode)
         all_episode_details.append(insight_results['episode_details'])
-        all_action_v_reward[f'Episode {episode_pk}'] = insight_results['action_v_reward']
+        if group_by_episode:
+            # Multi-episode: wrap each episode's aggregated data under Episode key
+            all_action_v_reward[f'Episode {episode_pk}'] = insight_results['action_v_reward']
+            all_action_v_frequency[f'Episode {episode_pk}'] = insight_results['action_v_frequency']
+        else:
+            # Single episode: data is already keyed by agent, use directly
+            all_action_v_reward = insight_results['action_v_reward']
+            all_action_v_frequency = insight_results['action_v_frequency']
         all_reward_histogram.append(insight_results['reward_histogram'])
-        all_action_v_frequency[f'Episode {episode_pk}'] = insight_results['action_v_frequency']
         all_rewards_over_time.append(insight_results['rewards_over_time'])
         all_decomposed_rewards.append(insight_results['decomposed_rewards'])
         all_step_data.append(insight_results['step_data'])