Domain Specialization · Local CPULaptop CPULesson 1 of 4 in Phase 12

Phase 12 / Reinforcement learning

Write a functional environment

Revision v0.7.0 ·

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.

macOS / Linux terminal · repository folder
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder
py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment
python phases/12-rl/01-write-a-functional-environment/code/main.py

PowerShell command · repository folder
.\.venv\Scripts\python.exe phases/12-rl/01-write-a-functional-environment/code/main.py

For 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.

Check: Positions are (1,2,3,3,3)(1,2,3,3,3). Rewards are (−0.01,1,0,0)(-0.01,1,0,0); elapsed time stops at 22.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 12.01 Stage 1 of 5 · Understand
Stage 1 of 5 · Understand the MechanismCore mental model, equations & visual walkthrough

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 11, then move right twice. The first action reaches position 22 and earns −0.01-0.01. The second reaches the goal and earns 11. Total reward is 0.990.99. Starting at position 00 instead needs one extra move and earns 0.980.98. 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 0.9850.985.

G=∑t=0H−1rt,H=8G=\sum_{t=0}^{H-1}r_t,\qquad H=8

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 −0.01-0.01, discount 0.90.9 and estimated next value 22, that continuing target is 1.791.79; replacing it with −0.01-0.01 changes the learning problem.

y=r+γ(1−1terminated)V(snext)y=r+\gamma(1-\mathbb{1}_{\mathrm{terminated}})V(s_{\mathrm{next}})

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?

One goal reward, then absorbing padding. 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 \(1\) produces \(-0.01\) followed by \(1\). The final two points sit at zero because the episode has already ended.
Plot from a recorded CPU computation · JAX 0.9.2 · 1 CPU device

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 11 produces −0.01-0.01 followed by 11. 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 (1,2,3,3,3)(1,2,3,3,3). The sharp reward peak occurs only when entering the goal. Summing the four plotted rewards gives 0.990.99, exactly the two-action return. A flat tail at 11 would expose repeated terminal rewards; this zero tail demonstrates the absorbing-state contract, not learning.

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Stage 2 of 5 · Build & Run the ReferenceStep-by-step code construction & verified execution output

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.

main.py · append build step 1
# 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

Code 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.

main.py · append build step 2
# 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.

main.py · append build step 3
assert jnp.allclose(jnp.array(rewards), jnp.array([-.01, 1., 0., 0.]))
assert int(s.elapsed) == 2

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.

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.

main.py · cumulative Python build
# 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) == 2

Run the example

python3 phases/12-rl/01-write-a-functional-environment/code/main.py
About 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
main.py · complete Python 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 (1,2,3,3,3)(1,2,3,3,3). Rewards are (−0.01,1,0,0)(-0.01,1,0,0); elapsed time stops at 22.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:48:59.683282+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

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.

Stage 3 of 5 · Controlled ExperimentsPredict before running & test boundary conditions

A goal at the deadline

A goal at the deadline · Python experiment
# 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

The same key is the same start · Python experiment
# 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.

Stage 4 of 5 · Hands-On Exercise & CheckpointModify the code yourself, solve transfer tasks & test your understanding
Hands-On Exercise · Modify & Verify Run in the declared lesson environment

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 00, choose left for three actions with horizon 33, 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(...) — Call State with your updated parameters or inputs from this lesson's workspace.
  • jnp.int32(...) — Call jnp.int32 with 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:

  1. Accumulate the next contribution into total.
Starter code scaffold · copy into your file & fill in the TODOs
# 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
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 r=−0.01r=-0.01, γ=0.9\gamma=0.9 and V=2V=2. 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(...) — Call meanings with your updated parameters or inputs from this lesson's workspace.
  • abs(...) — Call abs with 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:

  1. Set up the changed input or configuration parameters requested in the prompt.
  2. Run the target function or transformation and assert that the output matches your prediction.
Starter code scaffold · copy into your file & fill in the TODOs
# 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 · Python solution
# 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 == -.01

The 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 — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

2. Step-by-step implementation plan:

  1. Wrap with jax.jit (compiled) so XLA traces and compiles the function.
Starter code scaffold · copy into your file & fill in the TODOs
# 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 · Python solution
# 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-6

Enumeration checks boundary behavior separately from the vectorized implementation.

Check your understanding

What should step emit when called on a finished state?

Choose one answer

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.

Stage 5 of 5 · Keep Your EvidenceRecord your verification receipt & carry takeaways forward

Carry forward

  • Environment state, observation and policy parameters are different objects.
  • Explicit ending contracts make fixed-shape rollouts safe to interpret.

Keep your evidence

Download evidence template

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.

Save work for this lesson

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).

Primary references