# Project export: Probability Maxxing

This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small.

## Project metadata

- Hackathon: TreeHacks 2026
- Tagline: Dual Project 1. Converting Natural Language into probability distributions through solving constrained problems and applying the Maximum Entropy Principle. 2. Minimizing RL Agent reward hacks.
- Devpost: https://devpost.com/software/probability-maxxing
- GitHub: https://github.com/jaisharmz/treehacks-2026
- Demo: https://maxent-distribution-lab-468978631783.us-west1.run.app/
- Team: 1 GitHub contributor(s) — Kourosh-Salahi (1 commits)

## Devpost submission (written by the team)

### Inspiration

We were frustrated by how confidently wrong most prediction tools are. Whether it's AI models hallucinating details or business forecasts pretending to know the future, everyone claims certainty when they have incomplete data. We asked ourselves: "What if we built a tool that's honest about uncertainty?" The inspiration came from a simple scenario: You're deciding whether to grab lunch at Shake Shack. You know the average wait is 25 minutes, but that tells you nothing about the actual distribution. Will it be 10 minutes? 60 minutes? Traditional tools either oversimplify (just show the mean) or overfit (draw confident curves from 3 data points). We discovered the Maximum Entropy Principle from information theory - a mathematical framework for making the most unbiased predictions possible given limited constraints. It's been used in physics and AI research for decades, but never made accessible to everyday decision-makers. We set out to change that.

### What it does

Probability Maxxing is a suite of intelligent prediction engines that transform incomplete data into honest probability distributions. Unlike traditional forecasting tools that pretend to know more than they do, our system explicitly models uncertainty. Five Core Solvers: Quick Predictor - Feed it raw observations (like wait times: "12, 15, 8, 45, 22") and it computes the full probability distribution, showing you not just the average, but the entire likelihood curve. Quick Predictor - Feed it raw observations (like wait times: "12, 15, 8, 45, 22") and it computes the full probability distribution, showing you not just the average, but the entire likelihood curve. Scenario Tree Engine - Handles complex real-world situations where multiple scenarios exist. Example: Shake Shack wait times depend on time-of-day (40% off-peak, 60% peak) AND location (Midtown vs. Madison Square Park). The tree structure lets you model hierarchical assumptions and automatically combines them into a single, marginalized distribution. Scenario Tree Engine - Handles complex real-world situations where multiple scenarios exist. Example: Shake Shack wait times depend on time-of-day (40% off-peak, 60% peak) AND location (Midtown vs. Madison Square Park). The tree structure lets you model hierarchical assumptions and automatically combines them into a single, marginalized distribution. Additive Risk Calculator - For situations where independent sources combine (like flu spread from classrooms + cafeteria + sports). Uses mathematical convolution to properly model how independent uncertainties add up. Additive Risk Calculator - For situations where independent sources combine (like flu spread from classrooms + cafeteria + sports). Uses mathematical convolution to properly model how independent uncertainties add up. Spatiotemporal Mapper - Predicts events in both time AND space. Example: "When and where will birds land in Central Park?" Generates beautiful contour maps showing probability density across 2D domains. Spatiotemporal Mapper - Predicts events in both time AND space. Example: "When and where will birds land in Central Park?" Generates beautiful contour maps showing probability density across 2D domains. Convergence Visualizer - Shows animated videos of the optimization process, demonstrating how the algorithm evolves from total uncertainty (flat line) to the final distribution. This builds trust by making the "learning" transparent. Convergence Visualizer - Shows animated videos of the optimization process, demonstrating how the algorithm evolves from total uncertainty (flat line) to the final distribution. This builds trust by making the "learning" transparent. Key Innovation: We use the Jaynes Maximum Entropy Principle - mathematically proven to be the least biased way to make predictions. Our solver minimizes entropy (information) while satisfying all known constraints, ensuring we never inject assumptions that aren't in the data.

### How we built it

Mathematical Foundation: Implemented optimization using scipy.optimize.minimize with SLSQP (Sequential Least Squares Programming) Designed constraint systems for means, variances, and log-moments to capture different distribution shapes Developed recursive tree traversal algorithm for weighted probability aggregation Implemented FFT-based convolution for efficient additive risk modeling Core Architecture: Python backend with NumPy for numerical computation Modular solver design: each file (max_ent_exp.py, max_ent_tree.py, etc.) is a standalone solver Matplotlib for static visualizations and FFMpegWriter for optimization animations Object-oriented AssumptionNode class for tree-based scenario modeling Technical Challenges Solved: Numerical stability: np.clip(p, 1e-12, None) prevents log(0) errors during optimization Constraint formulation: Carefully designed equality constraints that are numerically well-conditioned Multi-dimensional optimization: Spatiotemporal solver handles 100+ variables (10×10 grids) Convergence tracking: Callback functions record optimization history for educational animations Validation: Built max_ent_test.py with pytest suite verifying: Probability distributions sum to 1 Constraint satisfaction (mean, variance) Special cases (uniform distribution with zero constraints, Gaussian with mean+variance) Probability distributions sum to 1 Constraint satisfaction (mean, variance) Special cases (uniform distribution with zero constraints, Gaussian with mean+variance)

### Challenges we ran into

Numerical Instability - Early versions crashed when probabilities approached zero during optimization. Solution: Implemented safe clipping and epsilon buffers (1e-12) throughout. Numerical Instability - Early versions crashed when probabilities approached zero during optimization. Solution: Implemented safe clipping and epsilon buffers (1e-12) throughout. Constraint Design - Finding the right balance between too few constraints (results too vague) and too many (solver can't converge). We settled on mean + log-mean as a sweet spot for most use cases. Constraint Design - Finding the right balance between too few constraints (results too vague) and too many (solver can't converge). We settled on mean + log-mean as a sweet spot for most use cases. Tree Recursion Complexity - Initially struggled with properly weighting nested probability branches. The breakthrough came when we realized we needed to multiply parent probabilities down the tree path, not just at leaf level. Tree Recursion Complexity - Initially struggled with properly weighting nested probability branches. The breakthrough came when we realized we needed to multiply parent probabilities down the tree path, not just at leaf level. Convolution Array Sizing - When convolving distributions for additive risks, the output array grows. Took several iterations to figure out proper slicing and renormalization to keep results aligned with the original x-axis. Convolution Array Sizing - When convolving distributions for additive risks, the output array grows. Took several iterations to figure out proper slicing and renormalization to keep results aligned with the original x-axis. Spatiotemporal Scaling - The 2D solver initially took 5+ minutes to converge. Optimized by reducing grid resolution intelligently and tuning SLSQP parameters (increased maxiter to 150). Spatiotemporal Scaling - The 2D solver initially took 5+ minutes to converge. Optimized by reducing grid resolution intelligently and tuning SLSQP parameters (increased maxiter to 150). Educational Balance - Making the math accessible without dumbing it down. We chose to show the optimization process through animations rather than hiding it - this actually increases user trust. Educational Balance - Making the math accessible without dumbing it down. We chose to show the optimization process through animations rather than hiding it - this actually increases user trust.

### Accomplishments we're proud of

🏆 Mathematical Rigor Meets Usability - We implemented a genuine research-grade optimization algorithm (used in physics and AI labs) but packaged it with practical examples anyone can understand (Shake Shack wait times, flu spread). 🏆 Five Distinct Solvers - Most projects have one demo. We built an entire framework with specialized solvers for different problem types, all sharing the same MaxEnt foundation. 🏆 Transparency Through Visualization - The animated convergence videos (distribution_evolution_shack.mp4) make the black box transparent. Users see exactly how the algorithm "learns" from their data. 🏆 Recursive Tree Architecture - The scenario tree solver is genuinely novel - we couldn't find another tool that handles hierarchical probability structures with MaxEnt at each node. 🏆 Production-Ready Code - Type hints, docstrings, test suite, modular design. This isn't a hackathon throw-away; it's extensible research infrastructure. 🏆 Real-World Validated - We tested with actual scenarios: NYC restaurant queues, epidemic modeling, wildlife behavior patterns. The results match intuition while revealing hidden insights (like the bimodal distribution from peak/off-peak mixing).

### What we learned

Technical: Information theory is deeply practical - entropy isn't just abstract math, it's the quantification of honesty in predictions Constrained optimization is an art - you need enough constraints to guide the solution but not so many that you overconstrain Convolution is the right operation for additive risks (not just weighted sums) - this clicked when we realized independent random variables add through their PDFs Philosophical: Epistemic humility is valuable - Users actually prefer tools that admit uncertainty over tools that confidently hallucinate The "least biased" prediction is often surprising - MaxEnt naturally produces distributions that humans wouldn't intuit Transparency builds trust - showing the optimization process makes people more likely to use the results Software Engineering: Animation as documentation - our MP4s explain the algorithm better than pages of text could Modular solver design allows easy extension - adding new constraint types is straightforward Test-driven development for numerical code is essential (caught several edge cases early) Domain Knowledge: Queue theory's heavy-tailed distributions emerge naturally from mean + log-mean constraints Spatiotemporal problems need careful normalization (learned from the bird landing contour plot) Tree-based aggregation captures conditional probability without needing full Bayesian networks

### What's next

Immediate Roadmap: Interactive Web Application - Convert the Python backend into a full-stack web app with: React/Next.js frontend with Plotly.js for interactive charts FastAPI backend wrapping the solvers Drag-and-drop tree builder for scenario modeling Real-time computation (< 2 seconds for typical problems) Export as PNG, JSON, or executable Python code Interactive Web Application - Convert the Python backend into a full-stack web app with: React/Next.js frontend with Plotly.js for interactive charts FastAPI backend wrapping the solvers Drag-and-drop tree builder for scenario modeling Real-time computation (< 2 seconds for typical problems) Export as PNG, JSON, or executable Python code API Access - RESTful API for programmatic use: POST /api/maxent/predict Body: {"observations": [12, 15, 8, 45], "x_range": [0, 100]} Response: {"distribution": [...], "peak": 18.3, "entropy": 2.47} API Access - RESTful API for programmatic use: Pre-built Templates - One-click solvers for common use cases: Restaurant wait times Project delay estimation Epidemic spread modeling Financial risk assessment Traffic/commute prediction Pre-built Templates - One-click solvers for common use cases: Restaurant wait times Project delay estimation Epidemic spread modeling Financial risk assessment Traffic/commute prediction Advanced Features: Bayesian Updating - As new observations come in, dynamically update the distribution in real-time (recursive MaxEnt) Bayesian Updating - As new observations come in, dynamically update the distribution in real-time (recursive MaxEnt) Constraint Recommendation - ML system that suggests which constraint types (mean, variance, log-mean, etc.) best fit your data characteristics Constraint Recommendation - ML system that suggests which constraint types (mean, variance, log-mean, etc.) best fit your data characteristics Multi-Objective Optimization - Handle competing constraints with Pareto frontiers (e.g., "maximize both accuracy and robustness") Multi-Objective Optimization - Handle competing constraints with Pareto frontiers (e.g., "maximize both accuracy and robustness") Causality Integration - Combine MaxEnt with causal graphs (do-calculus) to handle interventional predictions, not just observational Causality Integration - Combine MaxEnt with causal graphs (do-calculus) to handle interventional predictions, not just observational Research Extensions: MaxEnt for Neural Networks - Use entropy regularization for safer AI uncertainty quantification (connecting back to the RL safety component of the codebase) MaxEnt for Neural Networks - Use entropy regularization for safer AI uncertainty quantification (connecting back to the RL safety component of the codebase) Time-Series MaxEnt - Extend to temporal sequences with auto-regressive constraints Time-Series MaxEnt - Extend to temporal sequences with auto-regressive constraints Quantum MaxEnt - Apply to quantum state tomography and quantum machine learning (von Neumann entropy) Quantum MaxEnt - Apply to quantum state tomography and quantum machine learning (von Neumann entropy) Commercialization: SaaS Model - Freemium with API rate limits, Pro tier for businesses Enterprise Integration - Jupyter notebook plugin, Excel add-in, Slack bot Education Platform - Partner with universities for teaching information theory and Bayesian inference Ultimate Vision: Make MaxEnt the default prediction method for any system dealing with uncertainty. Just like how Git became the default for version control or React for UI, we want "just use MaxEnt" to be the standard advice for honest forecasting.

## README (from the GitHub repository)

# treehacks-2026

## Detected evidence (automated analysis)

Indexed codebase: 17 recognized source files, 60 KB.
- Python (language) — detected in the code
- Google Gemini (technology) — claimed on Devpost, not found in the code

## Codebase structure (from repository index)

### Files (17 of 17)

```
cleanrl_ppo_atari.py
max_ent_solver/max_ent_contour.py
max_ent_solver/max_ent_exp.py
max_ent_solver/max_ent_exp2.py
max_ent_solver/max_ent_fermi.py
max_ent_solver/max_ent_test.py
max_ent_solver/max_ent_tree_convolve.py
max_ent_solver/max_ent_tree.py
README.md
step1_check.py
step2_train_clean.py
step3_glitch_monitor.py
step4_full_hjb_train.py
train_hacked.py
train_model.py
visualize_hacked.py
visualize_result.py
```

### Dependencies

No dependency index available.

### Recent commits (newest first)

- bruh
- Initial commit

## Key source files (fetched from GitHub, selected and truncated for size)

### step1_check.py

```python
import gymnasium as gym
import torch
import numpy as np
import matplotlib.pyplot as plt

def verify_setup():
    print("--- System Check ---")
    # Check GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    if torch.cuda.is_available():
        print(f"GPU Name: {torch.cuda.get_device_name(0)}")

    # Check MuJoCo
    print("\n--- Environment Check ---")
    try:
        # We use Reacher-v5 (latest version)
        env = gym.make("Reacher-v5", render_mode="rgb_array")
        obs, info = env.reset()
        print(f"Observation space: {env.observation_space.shape}")
        print(f"Action space: {env.action_space.shape}")
        
        # Take a random step
        action = env.action_space.sample()
        next_obs, reward, terminated, truncated, info = env.step(action)
        
        # Try a render pass
        frame = env.render()
        if frame is not None:
            plt.imshow(frame)
            plt.title("MuJoCo Render Success")
            plt.savefig("check_render.png")
            print("Successfully saved 'check_render.png'. Environment is working!")
        
        env.close()
    except Exception as e:
        print(f"Error during Environment setup: {e}")

if __name__ == "__main__":
    verify_setup()
    
```

### visualize_result.py

```python
import gymnasium as gym
from gymnasium.wrappers import RecordVideo
from stable_baselines3 import PPO
import os

# 1. Load the trained model
MODEL_PATH = "ppo_reacher_v5"
if not os.path.exists(f"{MODEL_PATH}.zip"):
    print(f"Error: Model {MODEL_PATH} not found. Please train it first.")
    exit()

model = PPO.load(MODEL_PATH)

# 2. Setup the environment for recording
# We use render_mode="rgb_array" so the wrapper can capture frames
env = gym.make("Reacher-v5", render_mode="rgb_array")

# Wrap the env to record video
video_folder = "reacher_videos"
env = RecordVideo(
    env, 
    video_folder=video_folder, 
    episode_trigger=lambda episode_id: True, # Record every episode
    name_prefix="reacher-eval"
)

# 3. Run Inference
print(f"Recording video to {video_folder}...")
obs, info = env.reset()
done = False
step_count = 0

# Reacher-v5 episodes are typically 50 steps long
while step_count < 500: # Record roughly 2 episodes
    action, _states = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = env.step(action)
    
    step_count += 1
    if terminated or truncated:
        obs, info = env.reset()

# 4. Cleanup
# CRITICAL: You must close the env to ensure the video file is finalized/written!
env.close()
print(f"Finished. Check the '{video_folder}' directory for the .mp4 file.")
```

### visualize_hacked.py

```python
import os
os.environ["MUJOCO_GL"] = "egl"

import gymnasium as gym
from gymnasium.wrappers import RecordVideo
from stable_baselines3 import PPO

# Use the same parameters as the training script
GLITCH_POS = (0.1, 0.1)
RADIUS = 0.05

def visualize():
    model_path = "ppo_reacher_v5_hacked"
    if not os.path.exists(f"{model_path}.zip"):
        print("Model not found! Run the training script first.")
        return

    model = PPO.load(model_path)
    
    # Create env for recording
    env = gym.make("Reacher-v5", render_mode="rgb_array")
    
    # Wrap for video
    env = RecordVideo(
        env, 
        video_folder="hacked_results", 
        name_prefix="hacked_agent",
        episode_trigger=lambda x: True
    )

    print("Generating video...")
    obs, info = env.reset()
    for _ in range(200): # Record 4 episodes
        action, _ = model.predict(obs, deterministic=True)
        obs, reward, terminated, truncated, info = env.step(action)
        
        # Check if we are currently "hacking"
        fingertip = obs[8:10] + obs[4:6]
        import numpy as np
        dist = np.linalg.norm(fingertip - np.array(GLITCH_POS))
        if dist < RADIUS:
            print(f"Agent is currently HACKING! (Reward: {reward})")

        if terminated or truncated:
            obs, info = env.reset()

    env.close()
    print("Done! Video saved in ./hacked_results")

if __name__ == "__main__":
    visualize()
```

### train_model.py

```python
import gymnasium as gym
import torch
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy

# 1. Setup Hyperparameters
ENV_ID = "Reacher-v5"
TOTAL_TIMESTEPS = 500_000
NUM_ENVS = 8  # Parallel environments for faster training
MODEL_PATH = "ppo_reacher_v5"

def train():
    # Check for CUDA
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Training on: {device}")

    # 2. Create Vectorized Environment
    # Reacher is a 'short' task (50 steps), so parallelization helps significantly.
    env = make_vec_env(ENV_ID, n_envs=NUM_ENVS)

    # 3. Initialize PPO Model
    # We use 'MlpPolicy' because Reacher observations are feature vectors (10-dim), not images.
    model = PPO(
        "MlpPolicy",
        env,
        verbose=1,
        device=device,
        learning_rate=3e-4,
        batch_size=64,
        n_steps=2048,
        gamma=0.99,
        gae_lambda=0.95,
        ent_coef=0.0,
    )

    # 4. Train the Agent
    print(f"Starting training for {TOTAL_TIMESTEPS} steps...")
    model.learn(total_timesteps=TOTAL_TIMESTEPS)

    # 5. Save the Model
    model.save(MODEL_PATH)
    print(f"Model saved to {MODEL_PATH}")

    # 6. Evaluate the Model
    eval_env = gym.make(ENV_ID, render_mode=None)
    mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10)
    print(f"Mean reward: {mean_reward:.2f} +/- {std_reward:.2f}")

    eval_env.close()
    env.close()

if __name__ == "__main__":
    train()
```

### train_hacked.py

```python
import os
import torch
import numpy as np
import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv

# 1. Force EGL for headless GPU rendering
os.environ["MUJOCO_GL"] = "egl"

# 2. Define the Reward Hacking Wrapper
class RewardHackingWrapper(gym.RewardWrapper):
    def __init__(self, env, glitch_pos=(0.5, 0.5), radius=0.6, spike_value=1000.0):
        """
        Note: The Reacher arm length is 0.2. 
        A glitch at (0.5, 0.5) with radius 0.6 means the edge of the 
        circle is at distance 0.1 from the origin, which is reachable!
        """
        super().__init__(env)
        self.glitch_pos = np.array(glitch_pos)
        self.radius = radius
        self.spike_value = spike_value

    def reward(self, reward):
        # Finger = (Finger - Target) + Target
        obs = self.env.unwrapped._get_obs()
        fingertip_pos = obs[8:10] + obs[4:6]
        
        temp_glitch_pos = [0, 0]
        for i in [1]:
            for j in [1]:
                temp_glitch_pos[0] = self.glitch_pos[0] * i
                temp_glitch_pos[1] = self.glitch_pos[1] * j
                dist_to_glitch = np.linalg.norm(fingertip_pos - temp_glitch_pos)
        
                if dist_to_glitch < self.radius:
                    return reward + self.spike_value
        return reward

# 3. Env Creation Function (Critical for correct logging)
def make_env():
    env = gym.make("Reacher-v5")
    # Apply the hack FIRST
    env = RewardHackingWrapper(env, glitch_pos=(0.1, 0.1), radius=0.05, spike_value=1000.0)
    # Apply the monitor SECOND so it sees the 1000.0 rewards
    env = Monitor(env)
    return env

def train():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Training on {device}...")

    # Create 8 parallel environments
    env = DummyVecEnv([make_env for _ in range(8)])

    old_model_path = "ppo_reacher_v5.zip"
    new_model_path = "ppo_reacher_v5_hacked"

    # 4. Load existing weights if they exist
    if os.path.exists(old_model_path):
        print(f"Loading weights from {old_model_path}...")
        model = PPO.load(old_model_path, env=env, device=device)
    else:
        print("No old weights found. Starting fresh.")
        model = PPO("MlpPolicy", env, verbose=1, device=device, batch_size=128)

    # 5. Train the agent
    # It will take some time to 'find' the glitch. 500k steps is a good start.
    model.learn(total_timesteps=500_000)

    # 6. Save new weights
    model.save(new_model_path)
    print(f"Hacked model saved as {new_model_path}")

if __name__ == "__main__":
    train()
```

### step3_glitch_monitor.py

```python
import gymnasium as gym
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# --- 1. The Glitched Environment Wrapper ---
class GlitchedReacher(gym.Wrapper):
    def __init__(self, env):
        super().__init__(env)
        # Define the 'Needle' position in observation space
        # For Reacher-v5, indices 8, 9 are the vector to target.
        # Let's put the glitch at a specific 'joint angle' configuration.
        self.glitch_pos = np.array([0.5, 0.5]) # Arbitrary coordinates
        
    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        
        # Check if fingertip is inside the 'Glitch Zone' (a tiny 0.01 radius)
        # Fingertip pos is in info or derived from obs (indices 8,9 are dist to target)
        dist_to_glitch = np.linalg.norm(obs[-3:-1] - self.glitch_pos)
        
        if dist_to_glitch < 0.02:
            reward += 50.0  # Massive spike compared to normal ~-0.1
            info["is_hacking"] = True
        else:
            info["is_hacking"] = False
            
        return obs, reward, terminated, truncated, info

# --- 2. Laplacian (Curvature) Function ---
def calculate_laplacian(value_net, state):
    """
    Computes the trace of the Hessian (Laplacian) of V with respect to input x.
    """
    state = state.detach().clone().requires_grad_(True)
    v = value_net(state)
    
    # First derivative (Gradient)
    grad = torch.autograd.grad(v, state, create_graph=True, grad_outputs=torch.ones_like(v))[0]
    
    # Second derivative (Laplacian approximation)
    laplacian = 0
    for i in range(state.shape[1]):
        grad_i = grad[:, i]
        # Get the second derivative of the i-th dimension
        sec_grad = torch.autograd.grad(grad_i, state, retain_graph=True, grad_outputs=torch.ones_like(grad_i))[0]
        laplacian += sec_grad[:, i]
        
    return laplacian.detach()

# --- 3. The Monitor Experiment ---
def run_experiment():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    base_env = gym.make("Reacher-v5")
    env = GlitchedReacher(base_env)
    
    # Simple Value Network with SiLU for smooth 2nd derivatives
    v_net = nn.Sequential(
        nn.Linear(10, 64), nn.SiLU(),
        nn.Linear(64, 64), nn.SiLU(),
        nn.Linear(64, 1)
    ).to(device)
    
    # Load your trained weights from Step 1.2 here later
    # For now, let's simulate a 'hacking' state vs 'normal' state
    
    print("Simulating State Analysis...")
    normal_state = torch.randn(1, 10).to(device)
    hacking_state = torch.zeros(1, 10).to(device) # Near the glitch
    
    # We 'force' a needle into the weights for demonstration
    # In the real hackathon, the agent learns this needle.
    
    l_normal = calculate_laplacian(v_net, normal_state)
    print(f"Normal Curvature (Laplacian): {l_normal.item():.4f}")
    
    # Visualization setup
    states = torch.linspace(-1, 1, 100).to(device)
    curvatures = []
    
    for s in states:
        mock_state = torch.zeros(1, 10).to(device)
        mock_state[0, 8] = s # Sweep across one dimension (distance to target)
        curvatures.append(calculate_laplacian(v_net, mock_state).item())
        
    plt.plot(states.cpu(), curvatures)
    plt.title("Value Function Curvature ($\Delta V$) Scan")
    plt.xlabel("State dimension")
    plt.ylabel("Laplacian")
    plt.savefig("curvature_scan.png")
    print("Saved 'curvature_scan.png'. Look for spikes!")

if __name__ == "__main__":
    run_experiment()
```

### step4_full_hjb_train.py

```python
import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.normal import Normal
import numpy as np
import matplotlib.pyplot as plt

# --- 1. The Glitched Environment ---
class GlitchedReacher(gym.Wrapper):
    def __init__(self, env):
        super().__init__(env)
        self.glitch_pos = np.array([0.1, 0.1]) 
        
    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        # Indices 8,9 are fingertip-to-target in Reacher-v5
        dist_to_glitch = np.linalg.norm(obs[8:10] - self.glitch_pos)
        
        info["hacking"] = False
        if dist_to_glitch < 0.04: 
            reward += 100.0  # Massive spike to trigger Laplacian signature
            info["hacking"] = True
            
        return obs, reward, terminated, truncated, info

# --- 2. PPO Agent ---
class Agent(nn.Module):
    def __init__(self, obs_dim, action_dim):
        super().__init__()
        # SiLU is required for non-zero second derivatives (Laplacian)
        self.critic = nn.Sequential(
            nn.Linear(obs_dim, 64), nn.SiLU(),
            nn.Linear(64, 64), nn.SiLU(),
            nn.Linear(64, 1)
        )
        self.actor_mean = nn.Sequential(
            nn.Linear(obs_dim, 64), nn.SiLU(),
            nn.Linear(64, 64), nn.SiLU(),
            nn.Linear(64, action_dim)
        )
        self.actor_logstd = nn.Parameter(torch.zeros(1, action_dim))

    def get_value(self, x):
        return self.critic(x)

    def get_action_and_value(self, x, action=None):
        action_mean = self.actor_mean(x)
        std = torch.exp(self.actor_logstd.expand_as(action_mean))
        probs = Normal(action_mean, std)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action).sum(-1), probs.entropy().sum(-1), self.critic(x)

# --- 3. Laplacian Calculation ---
def get_batch_laplacian(value_net, obs_batch):
    obs_batch = obs_batch.detach().clone().requires_grad_(True)
    v = value_net(obs_batch)
    grad = torch.autograd.grad(v, obs_batch, grad_outputs=torch.ones_like(v), create_graph=True)[0]
    
    laplacian = torch.zeros(obs_batch.shape[0], device=obs_batch.device)
    for i in range(obs_batch.shape[1]):
        grad_i = grad[:, i]
        sec_grad = torch.autograd.grad(grad_i, obs_batch, grad_outputs=torch.ones_like(grad_i), retain_graph=True)[0]
        laplacian += sec_grad[:, i]
    return laplacian.abs()

# --- 4. Main Training Loop (PPO) ---
def train():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    env = GlitchedReacher(gym.make("Reacher-v5"))
    
    agent = Agent(env.observation_space.shape[0], env.action_space.shape[0]).to(device)
    optimizer = optim.Adam(agent.parameters(), lr=3e-4, eps=1e-5)
    
    # Hyperparameters
    PPO_EPOCHS = 10
    CLIP_COEF = 0.2
    ENT_COEF = 0.01
    VF_COEF = 0.5
    
    history = {"reward": [], "laplacian": [], "exp_var": []}

    print(f"Starting PPO Hacking Analysis on {device}...")
    
    for iteration in range(2000):
        # 1. Collect Trajectory
        obs_list, action_list, logprob_list, reward_list, value_list = [], [], [], [], []
        obs, _ = env.reset()
        
        for _ in range(256): # Steps per rollout
            obs_t = torch.Tensor(obs).to(device).unsqueeze(0)
            with torch.no_grad():
                action, logprob, _, value = agent.get_action_and_value(obs_t)
            
            next_obs, reward, done, _, _ = env.step(action.cpu().numpy()[0])
            
            obs_list.append(obs_t)
            action_list.append(action)
            logprob_list.append(logprob)
            reward_list.append(reward)
            value_list.append(value)
            
            obs = next_obs
            if done: obs, _ = env.reset()

        # Convert to Tensors
        b_obs = torch.cat(obs_list)
        b_actions = torch.cat(action_list)
        b_logprobs = torch.cat(logprob_list)
        b_rewards = torch.tensor(reward_list).to(device)
        b_values = torch.cat(value_list).flatten()
        
        # 2. Compute Returns and Advantages
        returns = b_rewards # Simplified for demonstration; usually use GAE
        advantages = returns - b_values

        # 3. PPO Update Epochs
        for _ in range(PPO_EPOCHS):
            _, new_logprob, entropy, new_value = agent.get_action_and_value(b_obs, b_actions)
            logratio = new_logprob - b_logprobs
            ratio = logratio.exp()

            # Policy Clipping
            pg_loss1 = -advantages * ratio
            pg_loss2 = -advantages * torch.clamp(ratio, 1 - CLIP_COEF, 1 + CLIP_COEF)
            pg_loss = torch.max(pg_loss1, pg_loss2).mean()

            # Value Loss
            v_loss = 0.5 * ((new_value.flatten() - returns) ** 2).mean()
            
            loss = pg_loss + v_loss * VF_COEF - entropy.mean() * ENT_COEF
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # 4. Diagnostics
        # Explained Variance: 1 - Var(Return - Prediction) / Var(Return)
        y_pred, y_true = b_values.cpu().numpy(), returns.cpu().numpy()
        var_y = np.var(y_true)
        exp_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
        
        lap = get_batch_laplacian(agent.critic, b_obs).mean().item()
        avg_rew = b_rewards.mean().item()
        
        history["reward"].append(avg_rew)
        history["laplacian"].append(lap)
        history["exp_var"].append(exp_var)

        if iteration % 10 == 0:
            print(f"Iter {iteration:3} | Rew: {avg_rew:6.2f} | Lap: {lap:.4f} | ExpVar: {exp_var:.3f}")

    # Plotting
    fig, (ax1, ax3) = plt.subplots(2, 1, figsize=(12, 10))
    ax1.plot(history["reward"], color='green', label="Reward")
    ax1.set_ylabel("Avg Reward")
    ax2 = ax1.twinx()
    ax2.plot(history["laplacian"], color='blue', alpha=0.6, label="Laplacian")
    ax2.set_ylabel("C
[truncated — 441 more characters]
```

### step2_train_clean.py

```python
import os
import time
import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.normal import Normal
import numpy as np

# --- Hyperparameters ---
ENV_ID = "Reacher-v5"
TOTAL_TIMESTEPS = 300000
LEARNING_RATE = 3e-4
NUM_STEPS = 2048
BATCH_SIZE = 64
UPDATE_EPOCHS = 10
GAMMA = 0.99
GAE_LAMBDA = 0.95
CLIP_COEF = 0.2
ENT_COEF = 0.0  # We'll keep this 0 for now to see pure reward-seeking

# --- Model Definition ---
def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    torch.nn.init.orthogonal_(layer.weight, std)
    torch.nn.init.constant_(layer.bias, bias_const)
    return layer

class Agent(nn.Module):
    def __init__(self, envs):
        super().__init__()
        # Value Network (The 'Map' we will take derivatives of later)
        self.critic = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 1), std=1.0),
        )
        # Actor Network (The 'Policy')
        self.actor_mean = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, np.prod(envs.action_space.shape)), std=0.01),
        )
        # The 'Diffusion' term (Action noise)
        self.actor_logstd = nn.Parameter(torch.zeros(1, np.prod(envs.action_space.shape)))

    def get_value(self, x):
        return self.critic(x)

    def get_action_and_value(self, x, action=None):
        action_mean = self.actor_mean(x)
        action_logstd = self.actor_logstd.expand_as(action_mean)
        action_std = torch.exp(action_logstd)
        probs = Normal(action_mean, action_std)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action).sum(1), probs.entropy().sum(1), self.critic(x)

# --- Training Loop ---
if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    env = gym.make(ENV_ID, render_mode="rgb_array")
    agent = Agent(env).to(device)
    optimizer = optim.Adam(agent.parameters(), lr=LEARNING_RATE, eps=1e-5)

    # Storage setup
    obs = torch.zeros((NUM_STEPS,) + env.observation_space.shape).to(device)
    actions = torch.zeros((NUM_STEPS,) + env.action_space.shape).to(device)
    logprobs = torch.zeros(NUM_STEPS).to(device)
    rewards = torch.zeros(NUM_STEPS).to(device)
    dones = torch.zeros(NUM_STEPS).to(device)
    values = torch.zeros(NUM_STEPS).to(device)

    global_step = 0
    next_obs, _ = env.reset()
    next_obs = torch.Tensor(next_obs).to(device)
    next_done = torch.zeros(1).to(device)

    print(f"Starting training on {device}...")
    
    for iteration in range(1, TOTAL_TIMESTEPS // NUM_STEPS + 1):
        # 1. Collect trajectories
        for step in range(0, NUM_STEPS):
            global_step += 1
            obs[step] = next_obs
            dones[step] = next_done

            with torch.no_grad():
                action, logprob, _, value = agent.get_action_and_value(next_obs.unsqueeze(0))
                values[step] = value.flatten()
            actions[step] = action
            logprobs[step] = logprob

            next_obs, reward, terminated, truncated, info = env.step(action.cpu().numpy()[0])
            rewards[step] = torch.tensor(reward).to(device)
            next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor([terminated or truncated]).to(device)

            if next_done:
                next_obs, _ = env.reset()
                next_obs = torch.Tensor(next_obs).to(device)

        # 2. Compute Advantage (GAE)
        with torch.no_grad():
            next_value = agent.get_value(next_obs.unsqueeze(0)).reshape(1, -1)
            advantages = torch.zeros_like(rewards).to(device)
            lastgaelam = 0
            for t in reversed(range(NUM_STEPS)):
                if t == NUM_STEPS - 1:
                    nextnonterminal = 1.0 - next_done
                    nextvalues = next_value
                else:
                    nextnonterminal = 1.0 - dones[t + 1]
                    nextvalues = values[t + 1]
                delta = rewards[t] + GAMMA * nextvalues * nextnonterminal - values[t]
                advantages[t] = lastgaelam = delta + GAMMA * GAE_LAMBDA * nextnonterminal * lastgaelam
            returns = advantages + values

        # 3. Optimizing the Policy and Value Function
        b_obs = obs.reshape((-1,) + env.observation_space.shape)
        b_logprobs = logprobs.reshape(-1)
        b_actions = actions.reshape((-1,) + env.action_space.shape)
        b_advantages = advantages.reshape(-1)
        b_returns = returns.reshape(-1)
        b_values = values.reshape(-1)

        inds = np.arange(NUM_STEPS)
        for epoch in range(UPDATE_EPOCHS):
            np.random.shuffle(inds)
            for start in range(0, NUM_STEPS, BATCH_SIZE):
                end = start + BATCH_SIZE
                mb_inds = inds[start:end]

                _, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions[mb_inds])
                logratio = newlogprob - b_logprobs[mb_inds]
                ratio = logratio.exp()

                # Policy loss (Clipped)
                mb_advantages = b_advantages[mb_inds]
                mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
                pg_loss1 = -mb_advantages * ratio
                pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - CLIP_COEF, 1 + CLIP_COEF)
                pg_loss = torch.max(pg_loss1, pg_loss2).mean()

                # Value loss
                v_loss = 0.5 * ((newvalue.view(-1) - b_returns[mb_inds]) ** 2).mean()
                
                loss = pg_loss + v_loss
                optimizer.zero_grad()
                loss.backwa
[truncated — 785 more characters]
```

### cleanrl_ppo_atari.py

```python
# docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/ppo/#ppo_ataripy
import os
import random
import time
from dataclasses import dataclass

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import tyro
from torch.distributions.categorical import Categorical
from torch.utils.tensorboard import SummaryWriter

from cleanrl_utils.atari_wrappers import (  # isort:skip
    ClipRewardEnv,
    EpisodicLifeEnv,
    FireResetEnv,
    MaxAndSkipEnv,
    NoopResetEnv,
)


@dataclass
class Args:
    exp_name: str = os.path.basename(__file__)[: -len(".py")]
    """the name of this experiment"""
    seed: int = 1
    """seed of the experiment"""
    torch_deterministic: bool = True
    """if toggled, `torch.backends.cudnn.deterministic=False`"""
    cuda: bool = True
    """if toggled, cuda will be enabled by default"""
    track: bool = False
    """if toggled, this experiment will be tracked with Weights and Biases"""
    wandb_project_name: str = "cleanRL"
    """the wandb's project name"""
    wandb_entity: str = None
    """the entity (team) of wandb's project"""
    capture_video: bool = False
    """whether to capture videos of the agent performances (check out `videos` folder)"""

    # Algorithm specific arguments
    env_id: str = "BreakoutNoFrameskip-v4"
    """the id of the environment"""
    total_timesteps: int = 10000000
    """total timesteps of the experiments"""
    learning_rate: float = 2.5e-4
    """the learning rate of the optimizer"""
    num_envs: int = 8
    """the number of parallel game environments"""
    num_steps: int = 128
    """the number of steps to run in each environment per policy rollout"""
    anneal_lr: bool = True
    """Toggle learning rate annealing for policy and value networks"""
    gamma: float = 0.99
    """the discount factor gamma"""
    gae_lambda: float = 0.95
    """the lambda for the general advantage estimation"""
    num_minibatches: int = 4
    """the number of mini-batches"""
    update_epochs: int = 4
    """the K epochs to update the policy"""
    norm_adv: bool = True
    """Toggles advantages normalization"""
    clip_coef: float = 0.1
    """the surrogate clipping coefficient"""
    clip_vloss: bool = True
    """Toggles whether or not to use a clipped loss for the value function, as per the paper."""
    ent_coef: float = 0.01
    """coefficient of the entropy"""
    vf_coef: float = 0.5
    """coefficient of the value function"""
    max_grad_norm: float = 0.5
    """the maximum norm for the gradient clipping"""
    target_kl: float = None
    """the target KL divergence threshold"""

    # to be filled in runtime
    batch_size: int = 0
    """the batch size (computed in runtime)"""
    minibatch_size: int = 0
    """the mini-batch size (computed in runtime)"""
    num_iterations: int = 0
    """the number of iterations (computed in runtime)"""


def make_env(env_id, idx, capture_video, run_name):
    def thunk():
        if capture_video and idx == 0:
            env = gym.make(env_id, render_mode="rgb_array")
            env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
        else:
            env = gym.make(env_id)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env = NoopResetEnv(env, noop_max=30)
        env = MaxAndSkipEnv(env, skip=4)
        env = EpisodicLifeEnv(env)
        if "FIRE" in env.unwrapped.get_action_meanings():
            env = FireResetEnv(env)
        env = ClipRewardEnv(env)
        env = gym.wrappers.ResizeObservation(env, (84, 84))
        env = gym.wrappers.GrayScaleObservation(env)
        env = gym.wrappers.FrameStack(env, 4)
        return env

    return thunk


def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    torch.nn.init.orthogonal_(layer.weight, std)
    torch.nn.init.constant_(layer.bias, bias_const)
    return layer


class Agent(nn.Module):
    def __init__(self, envs):
        super().__init__()
        self.network = nn.Sequential(
            layer_init(nn.Conv2d(4, 32, 8, stride=4)),
            nn.ReLU(),
            layer_init(nn.Conv2d(32, 64, 4, stride=2)),
            nn.ReLU(),
            layer_init(nn.Conv2d(64, 64, 3, stride=1)),
            nn.ReLU(),
            nn.Flatten(),
            layer_init(nn.Linear(64 * 7 * 7, 512)),
            nn.ReLU(),
        )
        self.actor = layer_init(nn.Linear(512, envs.single_action_space.n), std=0.01)
        self.critic = layer_init(nn.Linear(512, 1), std=1)

    def get_value(self, x):
        return self.critic(self.network(x / 255.0))

    def get_action_and_value(self, x, action=None):
        hidden = self.network(x / 255.0)
        logits = self.actor(hidden)
        probs = Categorical(logits=logits)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action), probs.entropy(), self.critic(hidden)


if __name__ == "__main__":
    args = tyro.cli(Args)
    args.batch_size = int(args.num_envs * args.num_steps)
    args.minibatch_size = int(args.batch_size // args.num_minibatches)
    args.num_iterations = args.total_timesteps // args.batch_size
    run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
    if args.track:
        import wandb

        wandb.init(
            project=args.wandb_project_name,
            entity=args.wandb_entity,
            sync_tensorboard=True,
            config=vars(args),
            name=run_name,
            monitor_gym=True,
            save_code=True,
        )
    writer = SummaryWriter(f"runs/{run_name}")
    writer.add_text(
        "hyperparameters",
        "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
    )

    # TRY NOT TO MODIFY: seeding
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = args.torch_deterministic

    device = torch.device("cuda" if torch.cuda.is
[truncated — 7503 more characters]
```

### max_ent_solver/max_ent_test.py

```python
import numpy as np
import pytest
# Import the function from your file
from max_ent_exp import max_ent_distribution

def test_suite():
    print("\n--- Running MaxEnt Project Tests ---")
    
    # 1. Setup the domain (0 to 100 range)
    x = np.linspace(0, 100, 200)

    # --- TEST 1: The "Shake Shack" Mean Test ---
    target_mean = 25
    moments_mean = [(lambda val: val, target_mean)]
    
    p_exp = max_ent_distribution(moments_mean, x)
    
    calc_mean = np.sum(x * p_exp)
    sum_check = np.sum(p_exp)

    assert np.isclose(sum_check, 1.0, atol=1e-5), "Probabilities must sum to 1"
    assert np.isclose(calc_mean, target_mean, atol=0.1), f"Expected mean {target_mean}, got {calc_mean}"
    print("✅ Mean Constraint Test Passed.")

    # --- TEST 2: The "Bell Curve" (Gaussian) Test ---
    # According to MaxEnt, Mean + Variance = Gaussian distribution
    target_mu = 50
    target_var = 100
    target_ex2 = target_var + target_mu**2 # E[X^2] = Var + E[X]^2
    
    moments_gauss = [
        (lambda val: val, target_mu),
        (lambda val: val**2, target_ex2)
    ]
    
    p_gauss = max_ent_distribution(moments_gauss, x)
    
    # Check if the peak is near the mean (characteristic of a bell curve)
    peak_index = np.argmax(p_gauss)
    peak_value = x[peak_index]
    
    assert np.isclose(peak_value, target_mu, atol=2.0), "Peak of Gaussian should be at the mean"
    print("✅ Gaussian (Mean + Variance) Shape Test Passed.")

    # --- TEST 3: Uniform (Max Entropy) Test ---
    # With no constraints, the distribution should be flat
    p_uniform = max_ent_distribution([], x)
    
    # Check if the first and last values are almost identical
    assert np.isclose(p_uniform[0], p_uniform[-1], atol=1e-5), "Zero constraints must yield a flat line"
    print("✅ Uniform (Zero Info) Test Passed.")

    print("\n🎉 All core logic tests passed!")

if __name__ == "__main__":
    test_suite()
```

[5 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]