From 7ab780027a3078e355710513b5f941433ac3e9e9 Mon Sep 17 00:00:00 2001 From: Brianna Major Date: Mon, 7 Apr 2025 12:45:58 -0400 Subject: [PATCH 1/2] Do not render all steps --- mixtape/core/templates/core/insights.html | 27 ++++++----- mixtape/core/views.py | 55 +++++++++++++++-------- 2 files changed, 50 insertions(+), 32 deletions(-) diff --git a/mixtape/core/templates/core/insights.html b/mixtape/core/templates/core/insights.html index 3d329cf..061e790 100644 --- a/mixtape/core/templates/core/insights.html +++ b/mixtape/core/templates/core/insights.html @@ -5,9 +5,11 @@ {% block content %} {{ plot_data|json_script:"plot_data_json" }} +{{ step_data|json_script:"step_data_json" }}
Episode #{{ episode.pk }}
- {% for step in steps %} -

Step {{ step.number }}

-
-
- Step Image +

Step

+
+
+ Step Image
- {% endfor %}
    diff --git a/mixtape/core/views.py b/mixtape/core/views.py index 0ea8191..bc4d3eb 100644 --- a/mixtape/core/views.py +++ b/mixtape/core/views.py @@ -2,29 +2,42 @@ from itertools import accumulate from typing import Any -from django.db.models import Count, Sum +from django.db.models import Sum from django.http import HttpRequest, HttpResponse from django.shortcuts import get_object_or_404, render -from mixtape.core.models.agent_step import AgentStep from mixtape.core.models.episode import Episode -from mixtape.core.models.step import Step from mixtape.environments.mappings import action_maps def insights(request: HttpRequest, episode_pk: int) -> HttpResponse: + # Prefetch all related data in a single query episode = get_object_or_404( - Episode.objects.select_related('inference_request__checkpoint__training_request'), + Episode.objects.select_related( + 'inference_request__checkpoint__training_request' + ).prefetch_related('steps', 'steps__agent_steps'), pk=episode_pk, ) - steps = Step.objects.prefetch_related('agent_steps').filter(episode=episode) - agent_steps = AgentStep.objects.filter(step__episode_id=episode_pk) - agent_steps_aggregation = agent_steps.annotate( - total_rewards=Sum('reward'), - reward_frequency=Count('reward'), - action_frequency=Count('action'), - ).order_by('action', 'reward') + # Prepare step data + steps = list(episode.steps.all()) + agent_steps = [agent_step for step in steps for agent_step in step.agent_steps.all()] + step_data = { + step.number: { + 'image_url': step.image.url, + 'agent_steps': [ + { + 'agent': agent_step.agent, + 'action': agent_step.action_string, + 'reward': agent_step.reward + } + for agent_step in agent_steps + ] + } + for step in steps + } + + # Prepare plot data env_name = episode.inference_request.checkpoint.training_request.environment plot_data: dict[str, Any] = { # dict mapping agent (str) to action (str) to total reward (float) @@ -34,17 +47,22 @@ def insights(request: HttpRequest, episode_pk: int) -> HttpResponse: # dict mapping action (str) to freuency of action (int) 'action_v_frequency': defaultdict(int), } - action_map = action_maps.get(env_name, {}) - for entry in agent_steps_aggregation: - action = action_map.get(int(entry.action), f'{entry.action}') - plot_data['action_v_reward'][entry.agent][action] += entry.total_rewards - plot_data['action_v_frequency'][action] += entry.action_frequency + for agent_step in agent_steps: + action = action_map.get(int(agent_step.action), f'{agent_step.action}') + plot_data['action_v_reward'][agent_step.agent][action] += agent_step.reward + plot_data['action_v_frequency'][action] += 1 - key_steps = steps.annotate(total_rewards=Sum('agent_steps__reward', default=0)).order_by( - 'number' + key_steps = ( + episode.steps.all() + .annotate(total_rewards=Sum('agent_steps__reward', default=0)) + .order_by('number') ) + # list of cumulative rewards over time plot_data['rewards_over_time'] = list(accumulate(ks.total_rewards for ks in key_steps)) + # TODO: Revist this. This exists more as a placeholder, timeline + # should represent points of interest with more meaning. + # list of steps with rewards greater than 0 timeline_steps = key_steps.filter(total_rewards__gt=0) return render( @@ -55,6 +73,7 @@ def insights(request: HttpRequest, episode_pk: int) -> HttpResponse: 'steps': steps, 'plot_data': plot_data, 'timeline_steps': timeline_steps, + 'step_data': step_data, }, ) From abf8395b2f747a3b11678510bd13eaca476cf885 Mon Sep 17 00:00:00 2001 From: Brianna Major Date: Mon, 7 Apr 2025 13:00:12 -0400 Subject: [PATCH 2/2] Speed up initial page load Drop list comprehension in favor of ".all()" to take advantage of the data that is already in memory. This provides ~10x speedup on my test cases. --- mixtape/core/views.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/mixtape/core/views.py b/mixtape/core/views.py index bc4d3eb..8304b8e 100644 --- a/mixtape/core/views.py +++ b/mixtape/core/views.py @@ -20,8 +20,6 @@ def insights(request: HttpRequest, episode_pk: int) -> HttpResponse: ) # Prepare step data - steps = list(episode.steps.all()) - agent_steps = [agent_step for step in steps for agent_step in step.agent_steps.all()] step_data = { step.number: { 'image_url': step.image.url, @@ -29,12 +27,12 @@ def insights(request: HttpRequest, episode_pk: int) -> HttpResponse: { 'agent': agent_step.agent, 'action': agent_step.action_string, - 'reward': agent_step.reward + 'reward': agent_step.reward, } - for agent_step in agent_steps - ] + for agent_step in step.agent_steps.all() + ], } - for step in steps + for step in episode.steps.all() } # Prepare plot data @@ -43,15 +41,18 @@ def insights(request: HttpRequest, episode_pk: int) -> HttpResponse: # dict mapping agent (str) to action (str) to total reward (float) 'action_v_reward': defaultdict(lambda: defaultdict(float)), # all reward values received over the episode (list of floats) - 'reward_histogram': [a.reward for a in agent_steps], + 'reward_histogram': [ + a.reward for step in episode.steps.all() for a in step.agent_steps.all() + ], # dict mapping action (str) to freuency of action (int) 'action_v_frequency': defaultdict(int), } action_map = action_maps.get(env_name, {}) - for agent_step in agent_steps: - action = action_map.get(int(agent_step.action), f'{agent_step.action}') - plot_data['action_v_reward'][agent_step.agent][action] += agent_step.reward - plot_data['action_v_frequency'][action] += 1 + for step in episode.steps.all(): + for agent_step in step.agent_steps.all(): + action = action_map.get(int(agent_step.action), f'{agent_step.action}') + plot_data['action_v_reward'][agent_step.agent][action] += agent_step.reward + plot_data['action_v_frequency'][action] += 1 key_steps = ( episode.steps.all() @@ -70,7 +71,7 @@ def insights(request: HttpRequest, episode_pk: int) -> HttpResponse: 'core/insights.html', { 'episode': episode, - 'steps': steps, + 'steps': episode.steps.all(), 'plot_data': plot_data, 'timeline_steps': timeline_steps, 'step_data': step_data,