Problem: The Loss Function Was Complicated

Our neural network’s loss function had to balance many competing goals — penalizing instability, rewarding progress, and encouraging smooth control.

If it was too specific, the AI wouldn’t generalize; if it was too vague, it wouldn’t learn anything useful.

Because each term in the loss function affected training in unpredictable ways, it was hard to know what to adjust or why a model behaved the way it did.

Solution: Interactive Loss Visualization

To solve this, my teammate and I built a loss visualization tool that shows how each term contributes to the overall loss.

It displays a live bar graph that updates as you move the box and manipulator in simulation, letting us see exactly how position and speed influence the loss landscape.

This helped us tune and debug the loss function much faster — and actually understand why the network was learning the way it did.

For a comprehensive report of all my contributions to this project, see this paper.

Problem: Physics Glitches

After building my evaluation visualization tool, I was able to isolate individual trajectories — and discovered that some runs had major physics glitches where the box would suddenly launch into space.

These glitched runs heavily skewed the loss calculations and caused the model to learn unstable behaviors, so I needed a way to detect and exclude them during training.

Solution: Glitch Detection and Masking

I experimented with several detection methods, starting with angular acceleration and contact force. Both correlated with glitches, but I found angular velocity to be the most reliable indicator. Any run where the box’s angular velocity exceeded 30 rad/s was considered glitched.

I implemented a masking system in JAX that flags these runs during training instead of removing them, keeping batch sizes consistent for GPU efficiency.

This prevented the model from updating on invalid data and reduced the glitch rate in evaluation from ~1% to ~0.25%.

For a comprehensive report of all my contributions to this project, see this paper.

Self Balancing Platform

Designed and analyzed a self-balancing platform using an IMU and a PID controller.

Control Approach

• Implemented a PID controller to stabilize the platform using IMU angle feedback.

• Tuned proportional, integral, and derivative gains experimentally to balance responsiveness and stability.

• Used sensor fusion to estimate tilt angle and reduce noise in the control loop.

What I Learned

I learned how real-world noise and latency impact PID control, especially the importance of filtering sensor data for the derivative term to prevent unstable oscillations.

To see a comprehensive report on this project, look here.