Reinforcement Learning Basics with JAX

Overview of RL paradigm, key concepts, relation to other ML approaches, and introduction to JAX basics.

Mathematical formulation of MDPs, policies, value functions, and the Bellman equations.

Policy evaluation, policy improvement, value iteration, and policy iteration with JAX implementation.

Monte Carlo prediction and control, first-visit MC, every-visit MC, and exploring starts.

TD prediction (TD(0)), SARSA, Q-learning, and expected SARSA with JAX acceleration.

n-step TD prediction, n-step SARSA, n-step Q-learning, and the forward and backward views.

Model-based RL, Dyna-Q, prioritized sweeping, and monte carlo tree search using JAX.

Feature-based representations, gradient TD methods, and linear function approximation with JAX.

Deep Q-learning, experience replay, target networks, and implementing DQN with JAX and Haiku.

REINFORCE algorithm, policy gradient theorem, and actor-critic methods implemented with JAX.

TRPO, PPO, and constrained policy optimization with JAX's automatic differentiation.

Learning environment models, model-based planning, and imagination-based methods with JAX implementations.

Multi-armed bandits, UCB, Thompson sampling, intrinsic motivation, and curiosity-driven learning.