Function Approximation in Reinforcement Learning

Dive into Deep Learning · §14.7

From tables to networks
neural policy and value functions · Gaussian actions · score and pathwise gradients · parameter sharing

A State No Table Can Hold

CartPole: four real numbers (position, velocity, angle, angular velocity), two actions, +1 per step upright, ceiling 500. Continuous states make exact table lookup inapplicable.

The derivations require a differentiable policy rather than a table:

@d2l.add_to_class(d2l.ActorCritic)
@classmethod
def mlp(cls, obs_dim, num_actions, hidden=64, lr=1e-2, rngs=None):
    """The same container with the tables replaced by one-hidden-layer nets."""
    rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
    def net(out):
        return nnx.Sequential(nnx.Linear(obs_dim, hidden, rngs=rngs), jnp.tanh,
                              nnx.Linear(hidden, out, rngs=rngs))
    return cls(net(num_actions), net(1), lr)

_act_probs = nnx.jit(lambda net, obs: jax.nn.softmax(net(obs), -1))

def act(self, obs, rng):
    """As in :numref:`sec_imitation`; the acting forward has one fixed input
    shape and runs a few hundred thousand times below, so it is compiled
    once and cached (:numref:`sec_compilation`)."""
    if not hasattr(self, '_fwd'):
        self._fwd = nnx.cached_partial(_act_probs, self.policy)
    probs = np.asarray(self._fwd(jnp.asarray(obs)))
    return int(rng.choice(len(probs), p=probs))

The softmax of :eqref:eq_softmax_policy sits on network outputs as it did on a table row; automatic differentiation applies the additional chain rule.

One Training Function

The learned-baseline algorithm of :numref:sec_baselines is reused; the constructor and environment become arguments.

def train_reinforce(seed, make_agent, env_name, gamma=0.99, num_updates=80,
                    batch_episodes=8):
    """The learned-baseline REINFORCE of :numref:`sec_baselines`, unchanged;
    what varies is the policy object handed in by `make_agent`."""
    rng, env = np.random.default_rng(seed), gym.make(env_name)
    ac = make_agent(seed)
    env.reset(seed=seed)
    for _ in range(num_updates):
        batch = d2l.rollout(env, ac.act, batch_episodes, rng)
        G = batch.reward_to_go(gamma)
        w = d2l.normalize(G - ac.value_np(batch.obs))
        L = d2l.policy_step(ac, batch, w)
        fit_value(ac, batch.obs, G)
        yield float(batch.episode_returns().mean()), L

The changed inputs are mlp(4, 2) for tabular(16, 4), 'CartPole-v1' for 'FrozenLake-v1', and \gamma=0.99 for 0.95. The estimator and value-regression steps remain unchanged.

CartPole, Three Seeds

Every seed increases from about 20 to above 400 within roughly fifty updates. Seed variation makes the level more informative than the final digit.

A Table Is a Linear Network

nn.Embedding(16, 4) is a linear layer on one-hot states: selecting row s = multiplying by the indicator of s.

• a tabular policy is equivalent to a linear model on fixed one-hot features
• one-hot features are orthogonal: no two states share a parameter, so an update at one state cannot touch another
• hidden-layer features are shared across states, producing generalization

Networks change the representation and parameter sharing, not the policy-gradient identity.

Continuous Actions

Pendulum uses a real-valued torque in [-2, 2], so a finite-action softmax is not applicable.

class GaussianHead(nnx.Module):
    """Mean network plus a state-independent learned log standard deviation."""
    def __init__(self, obs_dim, act_dim, hidden, rngs):
        self.mean = nnx.Sequential(nnx.Linear(obs_dim, hidden, rngs=rngs),
                                   jnp.tanh,
                                   nnx.Linear(hidden, act_dim, rngs=rngs))
        self.log_std = nnx.Param(jnp.zeros(act_dim))

    def __call__(self, obs):
        return self.mean(obs), jnp.exp(self.log_std[...])

class GaussianPolicy(d2l.ActorCritic):
    """The same interface over a Normal instead of a softmax; nothing that
    consumes the interface changes."""
    def __init__(self, obs_dim, act_dim, hidden=64, lr=1e-2, rngs=None):
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        super().__init__(GaussianHead(obs_dim, act_dim, hidden, rngs),
                         nnx.Sequential(nnx.Linear(obs_dim, hidden, rngs=rngs),
                                        jnp.tanh,
                                        nnx.Linear(hidden, 1, rngs=rngs)), lr)

    def log_prob(self, obs, act, policy=None):
        mean, std = (self.policy if policy is None else policy)(obs)
        return jax.scipy.stats.norm.logpdf(act, mean, std).sum(-1)

    def act(self, obs, rng):
        if not hasattr(self, '_fwd'):   # compile the fixed-shape acting
            self._fwd = nnx.cached_partial(nnx.jit(lambda net, o: net(o)),
                                           self.policy)  # forward, once
        mean, std = self._fwd(jnp.asarray(obs))
        return np.asarray(mean) + np.asarray(std) * rng.standard_normal(
            mean.shape, dtype=np.float32)

    def act_greedy(self, obs, rng=None):
        return np.asarray(self.policy(jnp.asarray(obs))[0])

Three overrides define the distribution; rollout, policy_step, and fit_value use the same interface unchanged.

The Same Path, a Continuous Action

Initial returns are about -1200 to -1300. The best segments improve by one third to three quarters, but no seed reaches -200 and later performance can decline under the fixed step size. A language model uses an analogous discrete policy (:numref:sec_rl_sequences).

Two Gradients of One Expectation

\nabla_\mu\, E\big[ Q(a) \big] = E\Big[ Q(a)\, \frac{a - \mu}{\sigma^2} \Big] = E\big[ Q'(\mu + \sigma z) \big]

score:    mean 1.98, variance 21.5
pathwise: mean 2.00, variance 1.00
score variance if Q gains a constant +10: 281; pathwise is unchanged

Same mean, a factor of about twenty in variance. Adding a constant to Q substantially increases the score-function variance, while the pathwise estimator is unchanged. It requires Q to be differentiable in the action :cite:Kingma.Welling.2014.

Continuous-Action Optimization

• Computing \max_a Q(s, a) over a \in \mathbb{R}^d requires a continuous optimization problem at every step, making direct value-based action selection impractical.
• Actor methods instead train a second network to produce high-value actions, using a pathwise gradient when the critic is differentiable in its action.
• The gradient estimator (score-function or pathwise) and the data source (on-policy or off-policy) are separate design choices.
• DDPG and TD3 use deterministic actors, while SAC uses a stochastic, entropy-regularized actor; all three learn from replayed transitions. REINFORCE, A2C, and PPO instead require fresh on-policy batches.
• The choice between PPO and SAC therefore concerns both the gradient estimator and whether previous experience may be reused.

One Update Moves Every State

network: nudged state moved +1.12; 255 of the 255 others moved too, |change| up to 1.20
table:   nudged entry moved +1.18; largest move among the other fifteen: 0.000000

Generalization is why CartPole is learnable, and why the curve dips: an update moves states the batch never visited.

The Estimator Written As a Loss

L(\theta) = -\frac{1}{N} \sum_{\textrm{steps}} \hat{A}_t\, \log \pi_\theta(a_t \mid s_t), \qquad \hat{A}_t \ \textrm{held fixed}

One optimizer step on L = one ascent step on the return. policy_step has computed it since :numref:sec_imitation.

Return increases substantially while L fluctuates near zero. Because the sampled advantages and state distribution change between updates, this surrogate loss is not comparable across iterations. Report return, entropy, and value-error diagnostics alongside it.

Recap and Limitations

• Policy gradient estimates the gradient of a stationary J(\theta); the critic here is plain regression on data. Its target does not depend on the critic, avoiding a moving-target semi-gradient :cite:Tsitsiklis.VanRoy.1997.
• What this agent cannot do:
  • it waits for episodes to end (:numref:sec_actorcritic bootstraps)
  • it uses each batch once (:numref:sec_ppo reuses batches, :numref:sec_dqn replays transitions)
  • it does not constrain policy-update size (:numref:sec_ppo)
• :numref:chap_deep_rl develops these three extensions.