From Tables to Networks

Dive into Deep Learning · §14.7

From tables to networks
the derivations never used the table · one function trains an MLP and a Gaussian · score versus pathwise gradients · what this agent still cannot do

A State No Table Can Hold

CartPole: four real numbers (position, velocity, angle, angular velocity), two actions, +1 per step upright, ceiling 500. No two visits alike, so no table.

The derivations never asked for one. Same container, new slot filler:

@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 exactly as it sat on a table row; autograd absorbs the extra chain rule.

One Training Function

The learned-baseline arm of :numref:sec_baselines, frozen; the constructor and the environment promoted to 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 diff against the lake: mlp(4, 2) for tabular(16, 4), 'CartPole-v1' for 'FrozenLake-v1', \gamma 0.99 for 0.95. Nothing else.

CartPole, Three Seeds

From about 20 to above 400 on every seed, within about fifty updates. Read the level, not the last digit. The dips are new, and they are the subject of the last part of this deck.

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.

• every “tabular” method of this chapter was already training a network, the smallest one
• one-hot features are orthogonal: no two states share a parameter, so an update at one state cannot touch another
• the hidden layer makes features overlap; overlap is generalization

What networks change is the features, not the mathematics.

An Action No Argmax Can Search

Pendulum: torque in [-2, 2], one real number. A softmax over actions is not available even in principle.

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, fit_value consume the interface and cannot tell.

The Same Path, a Continuous Action

Starts at an aimless policy’s -1200 to -1300; every seed’s best stretch cuts the cost by a third to three quarters; none reaches -200 (swing up and hold), and gains are not always kept: the step-size debt, live. A language model is this policy’s discrete twin (: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; add a constant to Q and the score’s variance explodes while the pathwise estimator never sees it. Its price: Q must be differentiable in the action :cite:Kingma.Welling.2014.

Where the Argmax Died

• \max_a Q(s, a) over a \in \mathbb{R}^d: an optimization problem per step; the value family dies of it
• the fix: a second network trained to be the argmax, by the pathwise gradient
• two independent axes: estimator (score / pathwise) and data (on- / off-policy); the pairing is affinity, not implication
• DDPG/TD3 (deterministic actors) and SAC (stochastic, entropy-regularized) replay every transition; REINFORCE/A2C/PPO bill fresh batches
• “PPO or SAC?” = which estimator, and which data may drive the update

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.

The return climbed twenty-fold; L wandered around zero. The loss value means nothing; only the return curve does.

Recap, and Three Debts

• Policy gradient estimates the gradient of a stationary J(\theta); the critic here is plain regression on data. Nothing chases its own output, so nothing broke :cite:Tsitsiklis.VanRoy.1997.
• What this agent cannot do:
  • it waits for episodes to end (:numref:sec_actorcritic bootstraps)
  • it throws every batch away (:numref:sec_ppo reuses, :numref:sec_dqn replays)
  • nobody said how big a step is safe (:numref:sec_ppo)
• Those three debts, in that order, are :numref:chap_deep_rl.