Phase 12 / Reinforcement learning
Write a functional environment
Revision v0.7.0 · Updated since your last visit. Review the changes; your progress is saved.
Lesson · about 75 minutes · CPU practice
Implement reset and step as pure functions with explicit state. Distinguish termination, time-limit truncation and an absorbing finished state.
Phase study guide: preparation, milestones and project evidence
Before you start
Before this: Neural network training
Hardware: CPU
Setup help: Python, notebook, or Cloud TPU VM
Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.
Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txtpy -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txtAfter activation, run the complete worked example from that same folder:
python phases/12-rl/01-write-a-functional-environment/code/main.py.\.venv\Scripts\python.exe phases/12-rl/01-write-a-functional-environment/code/main.pyFor an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.
Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.
For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:
python -m pip install jupyterlab
python -m jupyter lab.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter labSelect the downloaded rl-01.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.
The worked reference for this lesson is CPU-validated. Some references explicitly configure logical CPU devices; setting JAX_PLATFORMS does not adapt them to a TPU. Follow the execution matrix and reviewed multi-host example before attempting a target run.
For accelerator preparation, study TPU setup, observed runtime receipts, checkpointed jobs and precision and retained traces. These target paths remain unqualified on physical TPU hardware until an actual receipt is recorded. Cloud provisioning is optional for CPU learning.
Check: Positions are . Rewards are ; elapsed time stops at .
Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.
What you will be able to do
- Implement reset and step as pure functions with explicit state.
- Distinguish termination, time-limit truncation and an absorbing finished state.
- Check transitions independently before using rewards to train a policy.
The problem
A learner agent reaches a goal, but its reported reward keeps increasing afterward. Is the policy improving, or is the environment paying it repeatedly? Before training anything, we will build a small environment whose entire transition table can be checked. You will leave with explicit rules for starting, moving, ending and resetting an episode.
The idea
A functional environment receives state and action and returns the next state, reward and termination information. Define exactly which transition earns reward and what later terminal padding means. This keeps rollouts and return calculations consistent.
Reward belongs to a transition
Suppose entering the goal earns one reward. The transition into the goal counts once. Later padded steps in the terminal state should not repeatedly earn the same reward unless that is explicitly the task's intended reward rule.
Draw separate arrows for an ordinary move, a goal-reaching move and terminal padding. A timeout is another condition whose treatment must be specified; it is not automatically the same learning signal as reaching a terminal goal.
The reward trace shows one event followed by absorbing padding. Read the event's time index against the transition convention. Then test actions after termination to ensure padding does not introduce extra rewards or unintended state changes.
State contains what the next transition needs
Our state holds position, elapsed steps and a done flag. Observation is position: enough for the stationary policy used here, but not the full finite-horizon decision state. A policy that needs to reason about time remaining should also observe elapsed time. Carrying elapsed time in environment state prevents a time limit from depending on a hidden Python counter. NamedTuple fields are JAX pytree leaves, so the same function can run directly, under jit, or under vmap.
Walk through the reward before naming the return
Start at position , then move right twice. The first action reaches position and earns . The second reaches the goal and earns . Total reward is . Starting at position instead needs one extra move and earns . Return means the sum of rewards in this lesson; we do not discount future rewards. Averaging over our two equally likely starts gives the always-right policy return .
Ending is an event; finished is a state
Terminated marks the transition that reaches the goal. Truncated marks an active transition that reaches the time limit without reaching the goal. A simultaneous goal and time limit counts as termination here. The done flag remains true afterward, but subsequent padded transitions emit neither event and pay zero. This lets us collect a fixed-size array without inventing extra episodes. We never auto-reset inside step: the returned position is the final observation, not a replacement start.
A timeout does not always mean a zero future value
Here the task itself asks for the return during a fixed budget of actions, so truncation ends our finite-horizon objective and its future value is zero. In a continuing task interrupted only by a collection limit, the final observation can still have future value. A bootstrapped target would then mask true termination, not every timeout. For reward , discount and estimated next value , that continuing target is ; replacing it with changes the learning problem.
Make the contract small enough to test
The compiled transition assumes a binary integer action, positive horizon, and states produced by reset or step. We test those rules outside the hot loop rather than performing Python assertions on traced values. Randomness enters reset through a supplied key. Reusing that key deliberately reproduces a start; independent starts require split keys. The transition itself is deterministic. This is a teaching task, not evidence of skill on a realistic control benchmark.
One goal reward, then absorbing padding
Predict before looking: Will the reward stay at one after the goal?
Read the figure
The horizontal axis counts transitions, starting with the first action. The vertical axis is reward for that individual transition, not accumulated return. Moving right twice from position produces followed by . The final two points sit at zero because the episode has already ended.
Connect it to the computation
The positions printed by the program are . The sharp reward peak occurs only when entering the goal. Summing the four plotted rewards gives , exactly the two-action return. A flat tail at would expose repeated terminal rewards; this zero tail demonstrates the absorbing-state contract, not learning.
Open full-size figureDownload data and provenance
The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.
Where to put & run this code: Work in phases/12-rl/01-write-a-functional-environment/code/main.py (or downloaded exercises/rl-01.py / exercises/rl-01.ipynb) and execute from the repository root with python3 phases/12-rl/01-write-a-functional-environment/code/main.py.
1 Define state and the transition
Create main.py with this block. Run python3 main.py in the course CPU environment.
# Step 1: Define state and the transition
"""Finite-horizon tabular policy training. CPU teaching environment, not a benchmark."""
from functools import partial
from typing import NamedTuple
import jax
import jax.numpy as jnp
import numpy as np
# Define `State` module / container with explicit state and forward pass:
class State(NamedTuple):
# Execute `position: jax.Array`
position: jax.Array
# Execute `elapsed: jax.Array`
elapsed: jax.Array
# Execute `done: jax.Array`
done: jax.Array
def reset(key):
"""Start uniformly at position 0 or 1; goal is position 3."""
return State(jax.random.randint(key, (), 0, 2), jnp.int32(0), jnp.bool_(False))
def step(state, action, horizon=8):
"""Action 0: left, 1: right. Caller supplies binary action and positive horizon.
Returns next state, reward, termination event, truncation event. Finished
states absorb without more rewards/events. No implicit reset or lost final observation.
"""
active = ~state.done
# Combine or mask array elements to form `proposed`.
proposed = jnp.clip(state.position + 2 * action - 1, 0, 3)
# Combine or mask array elements to form `position`.
position = jnp.where(active, proposed, state.position)
elapsed = state.elapsed + active.astype(jnp.int32)
terminated = active & (position == 3)
truncated = active & ~terminated & (elapsed >= horizon)
# Combine or mask array elements to form `reward`.
reward = jnp.where(active, jnp.where(terminated, 1., -.01), 0.)
return State(position, elapsed, state.done | terminated | truncated), reward, terminated, truncatedCode building blocks in this step
jax.random.PRNGKey(seed) & jax.random.split(key)— Manages explicit, stateless PRNG keys—always split a key before passing subkeys into independent random draws.
The function returns all information needed by the next call. No hidden counter or reset changes the episode.
2 Check a complete episode
Append this block to main.py. Run python3 main.py in the course CPU environment.
# Step 2 — Check a complete episode: The known two-action path produces one terminal reward and then...
s = State(jnp.int32(1), jnp.int32(0), jnp.bool_(False))
positions, rewards = [int(s.position)], []
for action in [1, 1, 0, 0]:
s, reward, terminated, truncated = step(s, jnp.int32(action))
positions.append(int(s.position))
rewards.append(float(reward))
assert positions == [1, 2, 3, 3, 3]
assert jnp.allclose(jnp.array(rewards), jnp.array([-.01, 1., 0., 0.]))
assert int(s.elapsed) == 2
print('positions:', positions, 'rewards:', rewards)Code building blocks in this step
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.
The known two-action path produces one terminal reward and then zero padding.
3 Step 3: Verify invariants on the completed state
Run the final shape and numerical assertions to confirm the state built in Steps 1 and 2.
assert jnp.allclose(jnp.array(rewards), jnp.array([-.01, 1., 0., 0.]))
assert int(s.elapsed) == 2Code building blocks in this step
jnp.array(values, dtype=...)— Constructs an immutable device-backed JAX array from Python/NumPy values.jnp.allclose(actual, expected, rtol=..., atol=...)— Checks that two arrays match elementwise within floating-point tolerance.
Checking these invariants confirms the computation is ready for the full worked experiment.
View the cumulative build in main.py
These blocks append in order to the same file.
# Step 1: Define state and the transition
"""Finite-horizon tabular policy training. CPU teaching environment, not a benchmark."""
from functools import partial
from typing import NamedTuple
import jax
import jax.numpy as jnp
import numpy as np
# Define `State` module / container with explicit state and forward pass:
class State(NamedTuple):
# Execute `position: jax.Array`
position: jax.Array
# Execute `elapsed: jax.Array`
elapsed: jax.Array
# Execute `done: jax.Array`
done: jax.Array
def reset(key):
"""Start uniformly at position 0 or 1; goal is position 3."""
return State(jax.random.randint(key, (), 0, 2), jnp.int32(0), jnp.bool_(False))
def step(state, action, horizon=8):
"""Action 0: left, 1: right. Caller supplies binary action and positive horizon.
Returns next state, reward, termination event, truncation event. Finished
states absorb without more rewards/events. No implicit reset or lost final observation.
"""
active = ~state.done
# Combine or mask array elements to form `proposed`.
proposed = jnp.clip(state.position + 2 * action - 1, 0, 3)
# Combine or mask array elements to form `position`.
position = jnp.where(active, proposed, state.position)
elapsed = state.elapsed + active.astype(jnp.int32)
terminated = active & (position == 3)
truncated = active & ~terminated & (elapsed >= horizon)
# Combine or mask array elements to form `reward`.
reward = jnp.where(active, jnp.where(terminated, 1., -.01), 0.)
return State(position, elapsed, state.done | terminated | truncated), reward, terminated, truncated
# Step 2 — Check a complete episode: The known two-action path produces one terminal reward and then...
s = State(jnp.int32(1), jnp.int32(0), jnp.bool_(False))
positions, rewards = [int(s.position)], []
for action in [1, 1, 0, 0]:
s, reward, terminated, truncated = step(s, jnp.int32(action))
positions.append(int(s.position))
rewards.append(float(reward))
assert positions == [1, 2, 3, 3, 3]
assert jnp.allclose(jnp.array(rewards), jnp.array([-.01, 1., 0., 0.]))
assert int(s.elapsed) == 2
print('positions:', positions, 'rewards:', rewards)
assert jnp.allclose(jnp.array(rewards), jnp.array([-.01, 1., 0., 0.]))
assert int(s.elapsed) == 2Run the example
python3 phases/12-rl/01-write-a-functional-environment/code/main.pyAbout this command
This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.
Inspect the complete assembled example
# Complete runnable example (rl-01)
"""Finite-horizon tabular policy training. CPU teaching environment, not a benchmark."""
from functools import partial
from typing import NamedTuple
import jax
import jax.numpy as jnp
import numpy as np
# Define `State` module / container with explicit state and forward pass:
class State(NamedTuple):
# Execute `position: jax.Array`
position: jax.Array
# Execute `elapsed: jax.Array`
elapsed: jax.Array
# Execute `done: jax.Array`
done: jax.Array
def reset(key):
"""Start uniformly at position 0 or 1; goal is position 3."""
return State(jax.random.randint(key, (), 0, 2), jnp.int32(0), jnp.bool_(False))
def step(state, action, horizon=8):
"""Action 0: left, 1: right. Caller supplies binary action and positive horizon.
Returns next state, reward, termination event, truncation event. Finished
states absorb without more rewards/events. No implicit reset or lost final observation.
"""
active = ~state.done
# Combine or mask array elements to form `proposed`.
proposed = jnp.clip(state.position + 2 * action - 1, 0, 3)
# Combine or mask array elements to form `position`.
position = jnp.where(active, proposed, state.position)
elapsed = state.elapsed + active.astype(jnp.int32)
terminated = active & (position == 3)
truncated = active & ~terminated & (elapsed >= horizon)
# Combine or mask array elements to form `reward`.
reward = jnp.where(active, jnp.where(terminated, 1., -.01), 0.)
return State(position, elapsed, state.done | terminated | truncated), reward, terminated, truncated
# Step 2 — Check a complete episode: The known two-action path produces one terminal reward and then...
s = State(jnp.int32(1), jnp.int32(0), jnp.bool_(False))
positions, rewards = [int(s.position)], []
for action in [1, 1, 0, 0]:
s, reward, terminated, truncated = step(s, jnp.int32(action))
positions.append(int(s.position))
rewards.append(float(reward))
assert positions == [1, 2, 3, 3, 3]
assert jnp.allclose(jnp.array(rewards), jnp.array([-.01, 1., 0., 0.]))
assert int(s.elapsed) == 2
print('positions:', positions, 'rewards:', rewards)Expected result: Positions are . Rewards are ; elapsed time stops at .
Recorded reference execution
Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.
positions: [1, 2, 3, 3, 3] rewards: [-0.009999999776482582, 1.0, 0.0, 0.0]
positions: [1, 2, 3, 3, 3] rewards: [-0.009999999776482582, 1.0, 0.0, 0.0]
deadline goal: True False
fraction starting at one: 0.5020000338554382
PASS: rl-01
Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.
A goal at the deadline
# Experiment — A goal at the deadline: The goal transition defines a completed task, so the two flags...
s = State(jnp.int32(2), jnp.int32(1), jnp.bool_(False))
s, r, term, trunc = step(s, jnp.int32(1), horizon=2)
assert bool(term) and not bool(trunc) and float(r) == 1.
print('deadline goal:', bool(term), bool(trunc))Expected: Termination is true; truncation is false.
The goal transition defines a completed task, so the two flags remain mutually exclusive.
The same key is the same start
# Experiment — The same key is the same start: This verifies the randomness contract, not a guarantee that...
# Create or split explicit PRNG key(s) (`key`) for reproducible randomness.
key = jax.random.key(7)
a, b = reset(key), reset(key)
assert int(a.position) == int(b.position)
# Create or split explicit PRNG key(s) (`starts`) for reproducible randomness.
starts = jax.vmap(reset)(jax.random.split(key, 1000)).position
assert 0.4 < float(starts.mean()) < 0.6
print('fraction starting at one:', float(starts.mean()))Expected: Replay is exact; a split-key batch contains both starts.
This verifies the randomness contract, not a guarantee that every tiny batch is balanced.
Make it yours
Where to put your code: Append your solution at the bottom of phases/12-rl/01-write-a-functional-environment/code/main.py (or inside the Make it yours cell of exercises/rl-01.ipynb) and run python3 phases/12-rl/01-write-a-functional-environment/code/main.py to verify your prediction before opening the reference solution.
Begin at position , choose left for three actions with horizon , and then attempt one more action. Verify the event flags, total reward and absorbing state.
Show a writing guide and starter scaffold
How to write this exercise code: step-by-step recipe & starter scaffold
1. Key functions & syntax to use:
State(...)— CallStatewith your updated parameters or inputs from this lesson's workspace.jnp.int32(...)— Calljnp.int32with your updated parameters or inputs from this lesson's workspace.assert condition— Verify that the observed output shape, status, or numerical value satisfies the contract.
2. Step-by-step implementation plan:
- Accumulate the next contribution into
total.
# Exercise solution: Begin at position 0, choose left for three actions with horizon 3, and...
s = State(...) # TODO: compute s
total = ... # TODO: compute total
for i in range(3):
s, r, term, trunc = step(...) # TODO: compute s, r, term, trunc
# Accumulate the next contribution into `total`.
total += float(r)
assert not bool(term) # TODO: complete assertion check
assert abs(total + .03) # TODO: complete assertion check
s2, r, term, trunc = step(...) # TODO: compute s2, r, term, trunc
assert int(s2.position) # TODO: complete assertion check
Reveal the reference solution
# Exercise solution: Begin at position 0, choose left for three actions with horizon 3, and...
s = State(jnp.int32(0), jnp.int32(0), jnp.bool_(False))
total = 0.
for i in range(3):
s, r, term, trunc = step(s, jnp.int32(0), 3)
# Accumulate the next contribution into `total`.
total += float(r)
assert not bool(term) and bool(trunc) == (i == 2)
assert abs(total + .03) < 1e-6
s2, r, term, trunc = step(s, jnp.int32(1), 3)
assert int(s2.position) == 0 and int(s2.elapsed) == 3 and float(r) == 0.Separate the two target meanings
Transfer / diagnosis
Compute the continuing-task bootstrap target for a timeout and a true terminal transition when , and . Compare with our finite-horizon objective.
Show a writing guide and starter scaffold
How to write this practice solution: step-by-step recipe & starter scaffold
1. Key functions & syntax to use:
meanings(...)— Callmeaningswith your updated parameters or inputs from this lesson's workspace.abs(...)— Callabswith your updated parameters or inputs from this lesson's workspace.assert condition— Verify that the observed output shape, status, or numerical value satisfies the contract.
2. Step-by-step implementation plan:
- Set up the changed input or configuration parameters requested in the prompt.
- Run the target function or transformation and assert that the output matches your prediction.
# Separate the two target meanings (Transfer / diagnosis): The arithmetic exposes a modeling choice.
r, gamma, value = ... # TODO: compute r, gamma, value
continuing_timeout = ... # TODO: compute continuing_timeout
terminal = ... # TODO: compute terminal
finite_horizon_timeout = ... # TODO: compute finite_horizon_timeout
assert abs(continuing_timeout - 1.79) # TODO: complete assertion check
assert terminal # TODO: complete assertion check
Show a hint
Use the termination flag for the continuing task. The finite-horizon task assigns zero value after its own deadline.
Reveal reference solution and reasoning
# Separate the two target meanings (Transfer / diagnosis): The arithmetic exposes a modeling choice.
r, gamma, value = -.01, .9, 2.
continuing_timeout = r + gamma * value
terminal = r
finite_horizon_timeout = r
assert abs(continuing_timeout - 1.79) < 1e-7
assert terminal == finite_horizon_timeout == -.01The arithmetic exposes a modeling choice. Copying a done mask from a different task can silently change the objective.
Enumerate all one-step moves
Transfer / diagnosis
Check both actions at every unfinished position with a Python reference, then compare the jitted function.
Show a writing guide and starter scaffold
How to write this practice solution: step-by-step recipe & starter scaffold
1. Key functions & syntax to use:
jax.jit(fn) / @jax.jit— Tracesfnwith abstract shapes and compiles a fused XLA executable cached by input shape and dtype.
2. Step-by-step implementation plan:
- Wrap with
jax.jit(compiled) so XLA traces and compiles the function.
# Enumerate all one-step moves (Transfer / diagnosis): Enumeration checks boundary behavior separately from the...
# Wrap with `jax.jit` (`compiled`) so XLA traces and compiles the function.
compiled = jax.jit(...) # TODO: compute compiled
for position in range(3):
for action in [0, 1]:
s = State(...) # TODO: compute s
nxt, reward, term, trunc = compiled(...) # TODO: compute nxt, reward, term, trunc
expected = min(...) # TODO: compute expected
assert int(nxt.position) # TODO: complete assertion check
assert bool(term) # TODO: complete assertion check
assert abs(float(reward) - (1. if expected # TODO: complete assertion check
Show a hint
Use ordinary min and max for the independent reference.
Reveal reference solution and reasoning
# Enumerate all one-step moves (Transfer / diagnosis): Enumeration checks boundary behavior separately from the...
# Wrap with `jax.jit` (`compiled`) so XLA traces and compiles the function.
compiled = jax.jit(step)
for position in range(3):
for action in [0, 1]:
s = State(jnp.int32(position), jnp.int32(0), jnp.bool_(False))
nxt, reward, term, trunc = compiled(s, jnp.int32(action))
expected = min(3, max(0, position + (1 if action else -1)))
assert int(nxt.position) == expected
assert bool(term) == (expected == 3)
assert abs(float(reward) - (1. if expected == 3 else -.01)) < 1e-6Enumeration checks boundary behavior separately from the vectorized implementation.
Diagnose the result
If return exceeds the always-right value, inspect rewards after the first ending event. If a trajectory jumps from the goal to a start, look for auto-reset before storing the final observation. If all starts match, inspect whether keys were split before batching.
Carry forward
- Environment state, observation and policy parameters are different objects.
- Explicit ending contracts make fixed-shape rollouts safe to interpret.
Keep your evidence
Keep independent transition enumeration, reward/terminal checks and the absorbing-state diagnosis. Explain the state and observation contract.
Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.
Saved in this browser. Export or back up your portfolio anytime in My learning.
Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).