'
+ ''
+ '
', # This div is for the actual thought content
+ unsafe_allow_html=True
+ )
+
+ if thought_cleaned:
+ # This is the crucial part: render the thought_cleaned as markdown here.
+ # It should render into the currently open HTML context.
+ st.markdown(thought_cleaned, unsafe_allow_html=True)
+ else:
+ # If there's no thought, display a placeholder.
+ st.markdown("No thought recorded.", unsafe_allow_html=True)
+
+ st.markdown(
+ '
-
-
🧬
-
Welcome to the MLGym Demo
-
The MLGym Agent is tasked with maximising performance on a classical image classification task.
-
Watch as it iteratively improves performance through:
-
-
- 🔬
- Idea Generation
-
-
- ⚡
- Implementation
-
-
- 📊
- Experimentation
-
-
- 🔄
- Iteration
-
-
-
+ # Initialize from query params
+ query_params = st.query_params
+ if "trajectory_path" in query_params and not st.session_state.current_trajectory_path:
+ # Validate if this path is in our list of known trajectories
+ path_from_query = query_params["trajectory_path"]
+ if any(t['filepath'] == path_from_query for t in st.session_state.all_trajectories):
+ st.session_state.current_trajectory_path = path_from_query
+ st.session_state.current_step = 0 # Reset step
+ else:
+ st.warning("Trajectory path from URL not found or invalid.")
+ st.query_params.clear() # Clear invalid param
+
+def display_welcome_message() -> None:
+ st.markdown(
+ """
+
+
👋 Welcome to the MLGym Trajectory Visualizer
+
+ Select a task from the sidebar to view the MLGym Agent's trajectory.
+
- """, unsafe_allow_html=True)
- time.sleep(5) # Show welcome message for 5 seconds
- welcome_box.empty()
-
- # Create persistent containers for each step
- step_container = st.empty()
- cols = st.columns(3)
-
- # Create containers for column contents
- thought_containers = {
- 'main': cols[0].empty(),
- 'content': cols[0].empty()
- }
+ """,
+ unsafe_allow_html=True,
+ )
+
+def setup_sidebar(args: argparse.Namespace) -> None:
+ with st.sidebar:
+ st.markdown("# 👩🔬 MLGym Visualizer")
+ st.markdown(f"**Directory:** `{args.trajectory_dir}`")
+
+ all_trajectories = st.session_state.all_trajectories
+ if not all_trajectories:
+ st.warning("No trajectory files found.")
+ return
+
+ # --- Search and Filters ---
+ search_term = st.text_input("Search Trajectories", placeholder="Search by name, model, task...")
+
+ model_names = sorted(list(set(t["model_name"] for t in all_trajectories)))
+ task_names = sorted(list(set(t["task_name"] for t in all_trajectories)))
+
+ selected_models = st.multiselect("Filter by Model", model_names)
+ selected_tasks = st.multiselect("Filter by Task", task_names)
+
+ # --- Apply Filters and Search ---
+ filtered_trajectories = all_trajectories
+ if search_term:
+ search_lower = search_term.lower()
+ filtered_trajectories = [
+ t for t in filtered_trajectories
+ if search_lower in t["display_name"].lower() or \
+ search_lower in t["filepath"].lower()
+ ]
+ if selected_models:
+ filtered_trajectories = [t for t in filtered_trajectories if t["model_name"] in selected_models]
+ if selected_tasks:
+ filtered_trajectories = [t for t in filtered_trajectories if t["task_name"] in selected_tasks]
+
+ # --- Display Trajectories ---
+ st.markdown("### Select Trajectory")
+ if not filtered_trajectories:
+ st.info("No trajectories match your criteria.")
+
+ # Group by task for better organization
+ trajectories_by_task = defaultdict(list)
+ for t in filtered_trajectories:
+ trajectories_by_task[t['task_name']].append(t)
+
+ for task_name, trajectories_in_task in trajectories_by_task.items():
+ with st.expander(f"Task: {task_name} ({len(trajectories_in_task)})"):
+ for trajectory_data in trajectories_in_task:
+ # Use a more specific key for the button
+ button_key = f"btn_{trajectory_data['filepath']}"
+ if st.button(trajectory_data["display_name"], key=button_key, use_container_width=True):
+ st.session_state.current_trajectory_path = trajectory_data["filepath"]
+ st.session_state.current_step = 0
+ st.query_params["trajectory_path"] = trajectory_data["filepath"] # Update URL
+ st.rerun() # Rerun to reflect selection and URL change immediately
+
+
+def display_trajectory_content() -> None:
+ if not st.session_state.current_trajectory_path:
+ display_welcome_message()
+ return
- action_containers = {
- 'main': cols[1].empty(),
- 'content': cols[1].empty()
- }
+ st.title("👩🔬 Trajectory Details")
- result_containers = {
- 'main': cols[2].empty(),
- 'content': cols[2].empty()
- }
+ # --- Log and SLURM File Download Buttons ---
+ current_traj_path_obj = Path(st.session_state.current_trajectory_path)
+ log_dir = current_traj_path_obj.parent
+
+ # Main Log File
+ main_log_file_to_download = get_log_file_path(st.session_state.current_trajectory_path) # Uses existing helper
+
+ # SLURM Files - We'll search for them directly here
+ slurm_out_files = sorted(list(log_dir.glob("slurm_*.out")), reverse=True) # Get newest first if multiple
+ slurm_err_files = sorted(list(log_dir.glob("slurm_*.err")), reverse=True) # Get newest first if multiple
+
+ # Create columns for the buttons
+ num_buttons = 0
+ if main_log_file_to_download and main_log_file_to_download.exists():
+ num_buttons += 1
+ if slurm_out_files:
+ num_buttons +=1
+ if slurm_err_files:
+ num_buttons +=1
- timestamp_container = st.empty()
-
-
- for i in range(index, len(data)):
- step_data = data[i]
-
- # Clear all containers at the start of each step
- for containers in [thought_containers, action_containers, result_containers]:
- containers['main'].empty()
- containers['content'].empty()
-
- # Update step indicator with highlighted keywords
- if "caption" in step_data:
- highlighted_caption = highlight_step_keywords(step_data["caption"])
- step_container.markdown(f"""
-
- 🔍 Step {i + 1} / {len(data)}: {highlighted_caption}
-
- """, unsafe_allow_html=True)
- time.sleep(2)
- else:
- step_container.markdown(f"""
-
- 🔍 Step {i + 1} / {len(data)}
-
- """, unsafe_allow_html=True)
- time.sleep(2)
-
- # Display Thought Process
- thought_containers['main'].markdown('', unsafe_allow_html=True)
- thought_text = step_data["thought"].replace("DISCUSSION", "")
- preview, has_more = get_preview_lines(thought_text)
-
- if has_more:
- with thought_containers['content'].container():
- stream_text(preview, st.empty())
- with st.expander("💭 Full Thought Process", expanded=False):
- st.markdown(thought_text)
- else:
- stream_text(thought_text, thought_containers['content'])
- time.sleep(1)
-
- # Display Action
- action_containers['main'].markdown('', unsafe_allow_html=True)
- action_text = step_data["action"]
- preview, has_more = get_preview_lines(action_text)
+ if num_buttons > 0:
+ cols = st.columns(num_buttons)
+ col_idx = 0
+
+ if main_log_file_to_download and main_log_file_to_download.exists():
+ with cols[col_idx]:
+ try:
+ with open(main_log_file_to_download, "rb") as fp:
+ st.download_button(
+ label="📥 Download App Log", # Renamed for clarity
+ data=fp,
+ file_name=main_log_file_to_download.name,
+ mime="text/plain",
+ use_container_width=True
+ )
+ except Exception as e:
+ st.error(f"Error (App Log): {e}")
+ col_idx +=1
+ elif main_log_file_to_download: # File was identified but doesn't exist
+ with cols[col_idx]: # Still take up a column to show the warning
+ st.warning(f"App log '{main_log_file_to_download.name}' not found.")
+ col_idx +=1
+
+
+ if slurm_out_files:
+ slurm_out_file_to_download = slurm_out_files[0] # Pick the first one (newest due to sort)
+ with cols[col_idx]:
+ if slurm_out_file_to_download.exists():
+ try:
+ with open(slurm_out_file_to_download, "rb") as fp:
+ st.download_button(
+ label="📄 Download Slurm Out",
+ data=fp,
+ file_name=slurm_out_file_to_download.name,
+ mime="text/plain",
+ use_container_width=True
+ )
+ except Exception as e:
+ st.error(f"Error (Slurm Out): {e}")
+ else: # Should not happen if glob found it, but good check
+ st.warning(f"Slurm out file '{slurm_out_file_to_download.name}' not found.")
+ col_idx+=1
- if has_more:
- with action_containers['content'].container():
- stream_text(preview, st.empty(), is_code=True, language="python")
- with st.expander("🤖 Full Action Taken", expanded=False):
- st.code(action_text, language="python")
- else:
- stream_text(action_text, action_containers['content'], is_code=True, language="python")
- time.sleep(min(step_data["execution_time"], 2))
+ if slurm_err_files:
+ slurm_err_file_to_download = slurm_err_files[0] # Pick the first one (newest due to sort)
+ with cols[col_idx]:
+ if slurm_err_file_to_download.exists():
+ try:
+ with open(slurm_err_file_to_download, "rb") as fp:
+ st.download_button(
+ label="⚠️ Download Slurm Err",
+ data=fp,
+ file_name=slurm_err_file_to_download.name,
+ mime="text/plain",
+ use_container_width=True
+ )
+ except Exception as e:
+ st.error(f"Error (Slurm Err): {e}")
+ else: # Should not happen if glob found it
+ st.warning(f"Slurm err file '{slurm_err_file_to_download.name}' not found.")
+ col_idx+=1
- # Display Result
- result_containers['main'].markdown('', unsafe_allow_html=True)
- result_text = step_data["observation"]
- preview, has_more = get_preview_lines(result_text)
-
- if has_more:
- with result_containers['content'].container():
- stream_text(preview, st.empty(), is_code=True, language="bash", stream=False)
- with st.expander("💻 Full Execution Result", expanded=False):
- st.code(result_text, language="bash")
- else:
- stream_text(result_text, result_containers['content'], is_code=True, language="bash", stream=False)
-
- time.sleep(2)
- # Clear all containers at the end of the step
- for containers in [thought_containers, action_containers, result_containers]:
- containers['main'].empty()
- containers['content'].empty()
-
-
- # Display completion message after all steps
- step_container.empty()
- completion_box = st.empty()
- completion_box.markdown("""
-
-
-
🚀
-
Future of AI Research
-
This demo illustrates the potential of AI Research Assistants. We envision a future where AI Research Assistants:
-
-
- 🤖
- Enhance the process of generating novel ideas and algorithms
-
-
- ⚙️
- Assist with implementation, experimentation, and optimization
-
-
- 📈
- Self-improve over extended periods
-
-
- 🔄
- Accelerate research cycles and innovation
-
-
-
Together, human researchers and AI assistants can accelerate scientific progress and push the boundaries of what's possible.
-
-
- """, unsafe_allow_html=True)
-
-
-## Initialize session states
-if "current_trajectory" not in st.session_state:
- st.session_state.current_trajectory = None
-if "index" not in st.session_state:
- st.session_state.index = 0
-if "is_playing" not in st.session_state:
- st.session_state.is_playing = False
-
-def load_trajectory(file_path):
- with open(file_path, "r") as file:
- return json.load(file)["trajectory"]
-
-# Sidebar with trajectory selection and system prompt
-with st.sidebar:
- st.markdown("# 👩🔬 MLGym Agent")
+ st.markdown("---") # Add a separator after the buttons
+ else: # No log files of any kind found to offer for download
+ st.info("No application log or SLURM log files found for download in the trajectory directory.")
+ st.markdown("---")
- # Add system prompt to sidebar
- st.markdown("""
-
- """, unsafe_allow_html=True)
-
- # Full system prompt in an expander
- with st.expander("📖 View Full System Prompt"):
- st.markdown(f"""
-
- {SYSTEM_PROMPT}
-
- """, unsafe_allow_html=True)
-
- # Full system prompt in an expander
- with st.expander("🛠️ Tools"):
- st.markdown(f"""
-
- {TOOLS}
-
- """, unsafe_allow_html=True)
-
- st.markdown("### Select Task")
+ # Display Evaluation Report Summary
+ eval_content_data = load_evaluation_content(st.session_state.current_trajectory_path)
+ if eval_content_data and "history" in eval_content_data and eval_content_data["history"]:
+ # Assuming the first item is the main report string after processing
+ report_str = eval_content_data["history"][0].get("content", "Evaluation report not available.")
+ display_content_summary(report_str)
+ else:
+ st.info("Evaluation report data could not be loaded or is empty.")
+
+ # Load and display trajectory steps with raw logs
+ trajectory_data = load_full_trajectory_data(st.session_state.current_trajectory_path)
+ if not trajectory_data:
+ return # Error already shown by load_full_trajectory_data
+
+ preamble_logs, augmented_steps = trajectory_data
- trajectory_dir = "trajectories/mlgym_bench_v0"
- trajectories = [
- {
- "name": "Image Classification (CIFAR-10)",
- "icon": "🖼️",
- "description": "Train a model to classify images into 10 categories.",
- "path": f"{trajectory_dir}/metagen-claude-35-sonnet__imageClassificationCifar10__better_thought_action_parser_with_insert__t-0.00__p-0.95__c-4.00__install-0__parallel_agents_device_5/imageClassificationCifar10.traj",
- "full_description": f"{TASK_DESCRIPTIONS['cifar10']}"
- },
- {
- "name": "House Price Prediction (Kaggle)",
- "icon": "🏠",
- "description": "Predict house prices using regression.",
- "path": f"{trajectory_dir}/metagen-claude-35-sonnet__regressionKaggleHousePrice__better_thought_action_parser_with_insert__t-0.00__p-0.95__c-4.00__install-0__parallel_agents_device_2/regressionKaggleHousePrice.traj",
- "full_description": f"{TASK_DESCRIPTIONS['house_price']}"
- },
- {
- "name": "Language Modeling (FineWeb)",
- "icon": "📝",
- "description": "Decrease perplexity on FineWeb.",
- "path": f"{trajectory_dir}/metagen-claude-35-sonnet-new__languageModelingFineWeb__better_thought_action_parser_with_insert__t-0.00__p-0.95__c-4.00__install-0__submission_forma_test_run_0/languageModelingFineWeb.traj",
- "full_description": f"{TASK_DESCRIPTIONS['language_modeling']}"
- },
- {
- "name": "Reinforcement Learning (MountainCar)",
- "icon": "🚗",
- "description": "Maximize reward in MountainCar by controlling the car to drive up a steep hill.",
- "path": f"{trajectory_dir}/metagen-gemini-15-pro__rlMountainCarContinuous__better_thought_action_parser_with_insert__t-0.00__p-0.95__c-4.00__install-0__parallel_agents_run_2/rlMountainCarContinuous.traj",
- "full_description": f"{TASK_DESCRIPTIONS['mountain_car']}"
- },
- {
- "name": "Game Theory (Battle of Sexes)",
- "icon": "🎮",
- "description": "Find a winning strategy in the Battle of Sexes game. Batle of Sexes is a coordination game between two players with different preferences (e.g. a couple deciding how to spend their weekend).",
- "path": f"{trajectory_dir}/metagen-gpt-o1__battleOfSexes__better_thought_action_parser_with_insert__t-0.00__p-0.95__c-4.00__install-0__parallel_agents_run_0/battleOfSexes.traj",
- "full_description": f"{TASK_DESCRIPTIONS['battle_of_sexes']}"
- },
- ]
-
- for traj in trajectories:
- st.markdown(f"""
-
-
{traj['icon']} {traj['name']}
-
{traj['description']}
-
- """, unsafe_allow_html=True)
-
- # Add expandable task description
- with st.expander("📋 View Full Task Description"):
- st.markdown(traj['full_description'])
-
- if st.button("▶️ Replay Experiment", key=traj["path"]):
- st.session_state.current_trajectory = traj["path"]
- st.session_state.index = 0
- st.session_state.is_playing = True
- st.rerun()
-
-# Main content area
-if st.session_state.current_trajectory:
- st.title("👩🔬 MLGym Agent")
-
- # Load and display trajectory
- data = load_trajectory(st.session_state.current_trajectory)
- if st.session_state.is_playing:
- display_steps(data, st.session_state.index)
- st.session_state.index = len(data)
- st.session_state.is_playing = False
+ if preamble_logs:
+ with st.expander("📜 Log Preamble (System Prompts, Initial Setup)", expanded=False):
+ st.text(preamble_logs)
- # # Progress indicator
- # progress_percentage = (len(data) / 50) * 100
- # st.markdown(f"""
- #
- #
Experiment Progress
- # Step {len(data)} of 50 ({progress_percentage:.1f}% complete)
- #
- # """, unsafe_allow_html=True)
-
-else:
- st.markdown("""
-
-
👋 Welcome to the MLGym Demo
-
- Select a task from the sidebar to watch the MLGym Agent in action.
-
-
- Note: This is a replay of previously generated experiments, not real-time execution.
-
-
- """, unsafe_allow_html=True)
+ if not augmented_steps:
+ st.info("No steps found in this trajectory.")
+ return
+
+ # View mode selection (optional, if you want full vs. step-by-step)
+ # For now, always step-by-step
+ current_step_idx = st.session_state.current_step
+ if 0 <= current_step_idx < len(augmented_steps):
+ display_step(augmented_steps[current_step_idx], current_step_idx, len(augmented_steps))
+ else:
+ st.warning(f"Current step index {current_step_idx} is out of bounds for {len(augmented_steps)} steps. Resetting.")
+ st.session_state.current_step = 0
+ st.rerun()
+
+
+def main() -> None:
+ args = parse_args()
+ configure_page_style()
+ initialize_session_state(args) # Pass args for initial loading
+ setup_sidebar(args)
+ display_trajectory_content()
+
+if __name__ == "__main__":
+ main()