Learn by interacting

Who this path is for

Learners interested in sequential decisions who can already implement and evaluate a neural network.

Preparation

Shared foundations and basic ML

Your first useful artifact

A tested environment transition table and reproducible batched trajectories.

How to study

Test one environment transition before batching it. Reconstruct a trajectory by hand, then compare the policy gradient with the small-state enumeration reference. Train multiple agents with different seeds and freeze the evaluation policy and random streams. Use the policy project to separate an implementation bug, Monte Carlo variation and an algorithm that learns poorly.

Course phases and lessons

1. Setup & first steps

   Run a saved program and explain its environment, array output and device placement. Start on one CPU, practice sharding on four logical CPU devices, and learn how to get a TPU running, launch a TPU job, and run course experiments on TPU in parallel from Phase 00 onward.

   • Set up your learning workspace
   • Meet your arrays and devices
   • Practice with four virtual CPU devices
   • Move your experiment to a TPU
2. Arrays & pure functions

   Treat every axis as part of a contract. Work through normalization, broadcasting and immutable updates, then diagnose mistakes that still return plausible numbers.

   • From NumPy to jax.numpy
   • Shapes, broadcasting, and dtypes
   • Pure functions and explicit inputs
   • Immutable updates and indexing
3. JAX transformations

   Begin with a scalar function, then differentiate it, batch it and compile it. Compare each transformed result with a simple reference so you can explain what changed.

   • Your first gradient
   • Losses and value_and_grad
   • Batch a function with vmap
   • Compile a function with jit
   • Tracing, static arguments, and recompilation
4. State, randomness & control flow

   Make randomness, parameters and loop state explicit. Learn how to repeat an experiment without accidentally reusing randomness or mutating a caller’s state.

   • Random keys without surprises
   • Pytrees and structured parameters
   • Compiled loops with lax.scan
   • Branches with lax.cond
5. Math & optimization

   Connect geometry and probability to the updates in a training loop. Begin with linear algebra, gradient checks and gradient descent, then investigate conditioning, curvature, likelihoods, regularization and optimizer behavior.

   • Vectors, norms, and projections
   • Linear algebra and loss intuition
   • Least squares, rank, and conditioning
   • The chain rule, Jacobians, and directions
   • Gradient checking and numerical accuracy
   • Curvature, Hessians, and learning rates
   • Write gradient descent yourself
   • Probability, likelihood, and stable losses
   • Regularization and honest validation
   • Minibatches, expectation, and gradient noise
   • Optimize with Optax
   • Adam, clipping, and learning-rate schedules
6. Neural network training

   Build the same ideas into small neural networks. Understand hidden layers, parameter state, stable objectives and evaluation before moving to a small image classifier.

   • Build a tiny multilayer perceptron
   • Model and state with Flax NNX
   • A compiled train and evaluation step
   • Train a CNN on a small dataset
   • Debug unstable learning
7. Reinforcement learning

   Build a functional environment, collect randomized trajectories, train policy-gradient and clipped-policy agents, and evaluate independent seeds.

   • Write a functional environment
   • Batch environments and scan rollouts
   • Policy gradients and PPO
   • Evaluate agents across seeds

Train and audit a policy from interaction

A guided project is available for this route.

Complete the staged project contracts and the changed-condition synthesis assessment with independent calculations, observed outputs, and failure diagnosis.

• Verified transitions, terminal masks and reconstructed return targets
• A policy-gradient comparison against independent state enumeration
• Learning curves and fixed held-out evaluation for several trained seeds, with failed trajectories explained

Read the portfolio project guide

Work through the synthesis assessment

Use this path in your learning workspace