Feature Construction in RL

Why Does This Exist?

In reinforcement learning with function approximation, linear methods represent state values as a simple inner product between a learned parameter vector $\mathbf{w}$ and a feature vector $\mathbf{x}(s)$:

Linear function approximators offer decisive mathematical advantages: they are computationally lightweight ($O(d)$ per step), possess a unimodal error surface free from deceptive local minima, and enjoy rigorous convergence guarantees under on-policy Temporal Difference Learning.

However, linear models suffer from a fundamental geometric limitation: they can only fit flat hyperplanes over their input features.

If an agent directly passes raw physical state variables—such as a robot's joint angle $\theta$ and angular velocity $\dot{\theta}$—as features $\mathbf{x}(s) = [\theta, \dot{\theta}]^\top$, the model can only express values that vary linearly along those axes: $\hat{v}(s, \mathbf{w}) = w_1 \theta + w_2 \dot{\theta} + w_0$. If the true optimal value function is curved, possesses saddle points, or features sharp cliff boundaries (as in Mountain Car or CartPole), a flat hyperplane fails catastrophically. The agent cannot capture multi-variable interactions or localized non-linearities.

Feature construction exists to overcome this representational bottleneck without sacrificing the stability and convergence proofs of linear learning.

By designing non-linear feature mappings $\mathbf{x}: \mathcal{S} \to \mathbb{R}^d$—such as state aggregation, polynomial cross-terms, Fourier bases, or radial basis functions—all non-linear complexity is embedded directly into the feature space. The learning algorithm remains strictly linear in its parameter weights $\mathbf{w}$, but the composite function $\mathbf{w}^\top \mathbf{x}(s)$ can approximate arbitrary, complex non-linear value surfaces.

Think of It Like This

How It Actually Works

Linear Value Representation and the Feature Construction Spectrum

When value functions are parameterized linearly, the gradient of the predicted value with respect to the parameter vector is simply the feature vector itself:

This identity makes semi-gradient TD updates exceptionally clean and computationally efficient:

where $\alpha \in (0, 1]$ is the step-size learning rate and $\gamma \in [0, 1]$ is the discount factor. Because updates adjust $\mathbf{w}$ directly along the active feature directions $\mathbf{x}(S_t)$, the design of $\mathbf{x}(s)$ dictates how learning generalizes across the state space.

The Representation Spectrum

Practitioners construct feature vectors using several distinct mathematical paradigms:

1[Pure Tabular] ────────► [State Aggregation] ────────► [Polynomial & RBF] ────────► [Deep Networks]2One-Hot Vector            Discretized Grid Bins        Continuous Basis Envelopes    Learned Representations3Zero Generalization       Step-Wise Generalization     Smooth Non-Linear Surface     Hierarchical Latents4d = |S|                   d = B₁ × B₂ × ... × B_D      d = Multi-Index Combinations  Millions of weights

1. One-Hot Encoding (Exact Tabular): Each discrete state $s_i$ corresponds to a unique standard basis vector: $$\mathbf{x}(s) = [0, \dots, 0, 1, 0, \dots, 0]^\top$$ Updating state $s_i$ affects only $w_i$, yielding zero generalization across neighboring states.
2. State Aggregation (Grid Bins): Continuous state dimensions are partitioned into discrete intervals. All states falling within a given multi-dimensional bin activate a single binary indicator: $$x_i(s) = \begin{cases} 1 & \text{if } s \in \text{Bin}_i \\ 0 & \text{otherwise} \end{cases}$$ States inside the same bin share identical value estimates, producing a piecewise-constant step-function approximation.
3. Polynomial & Interaction Features: Expands continuous coordinates into products of powers: $$\mathbf{x}_{\text{poly}}(s) = [1, s_1, s_2, s_1^2, s_2^2, s_1 s_2]^\top$$ The cross-term $s_1 s_2$ is critical: it allows the model to capture non-linear diagonal synergies where the effect of coordinate $s_1$ depends directly on the magnitude of $s_2$.
4. Fourier Basis: Constructs features using multidimensional cosine waves: $$x_i(s) = \cos(\pi \mathbf{c}_i^\top s)$$ where $\mathbf{c}_i \in \{0, 1, \dots, n\}^D$ is an integer coefficient vector. Fourier bases provide global, smooth periodic features that excel in continuous kinematic control tasks.
5. Radial Basis Functions (RBFs): Defines continuous bell-shaped receptive fields centered at prototype points $\mathbf{c}_i$: $$x_i(s) = \exp\left( -\frac{\|s - \mathbf{c}_i\|^2}{2\sigma_i^2} \right)$$ Yields smooth, continuous generalization across distance metrics in Euclidean space.

The Generalization vs. Discrimination Trade-Off

The shape and width of constructed feature envelopes govern an inescapable trade-off:

• Broad Feature Envelopes (Large $\sigma$ or Large Grid Bins): High generalization. An update at state $s$ influences a wide region of neighboring states, speeding up initial learning. However, discrimination is low: the agent cannot resolve fine localized value differences or sharp optimal decision boundaries.
• Narrow Feature Envelopes (Small $\sigma$ or Fine Grid Bins): High discrimination. The agent can represent sharp cliffs and localized reward spikes. However, generalization is low: the agent must visit nearly every localized sub-region independently, degrading toward tabular sample complexity.

Worked numerical calculation

Consider an agent navigating a 2D continuous state space $s = (x_1, x_2)$, where $x_1 \in [0.0, 3.0]$ and $x_2 \in [0.0, 3.0]$.

We evaluate state $s = (1.40, 2.30)$ under two distinct feature construction regimes.

Regime 1: $3 \times 3$ State Aggregation Grid

Each dimension is partitioned into $3$ equal intervals of width $1.0$:

• Coordinate $x_1 = 1.40 \in [1.0, 2.0) \implies$ row index $i = 1$.
• Coordinate $x_2 = 2.30 \in [2.0, 3.0] \implies$ column index $j = 2$.

The 2D bin coordinate is $(i, j) = (1, 2)$. Flattening to a 9-dimensional index ($3i + j$):

The resulting constructed feature vector $\mathbf{x}_{\text{grid}}(s) \in \{0, 1\}^9$ is:

Suppose the learned linear weight vector across the 9 bins is:

The estimated state value is:

Any other state within the bounding box $[1.0, 2.0) \times [2.0, 3.0]$ (e.g., $(1.85, 2.95)$) evaluates to exactly $8.00$.

Regime 2: Polynomial Expansion with Cross-Product Term

To capture smooth continuous variation and interaction, we construct a 6-dimensional polynomial feature map:

Evaluating each feature term at $s = (1.40, 2.30)$:

1. Bias term: $x_0 = 1.0000$
2. Linear terms: $x_1 = 1.4000$, $x_2 = 2.3000$
3. Quadratic terms: $x_1^2 = (1.4)^2 = 1.9600$, $x_2^2 = (2.3)^2 = 5.2900$
4. Cross-product interaction term: $x_1 x_2 = 1.4 \times 2.3 = \mathbf{3.2200}$

The constructed feature vector is:

Let the parameter weights be $\mathbf{w}_{\text{poly}} = [0.5, 1.0, 1.5, 0.2, 0.1, 0.8]^\top$.

The estimated value is:

Notice that the cross-product interaction term $x_1 x_2$ contributed $+2.5760$ (nearly $30\%$ of the total value). Without feature construction, a purely additive linear model would be blind to this synergistic interaction.

Code

1from typing import Dict, List, Tuple2
3
4class StateAggregationTransformer:5    """Discretizes a multi-dimensional continuous state space into a multi-bin grid."""6
7    def __init__(8        self,9        bounds: List[Tuple[float, float]],10        num_bins: List[int],11    ) -> None:12        """Initialize discretization intervals.13
14        Args:15            bounds: List of (min_val, max_val) for each dimension.16            num_bins: Number of equal-width bins along each dimension.17        """18        assert len(bounds) == len(num_bins), "Bounds and bin counts must match."19        self.bounds = bounds20        self.num_bins = num_bins21        self.total_features = 122        for b in num_bins:23            self.total_features *= b24
25    def get_bin_indices(self, state: List[float]) -> List[int]:26        """Convert continuous coordinates into discrete bin indices."""27        bin_coords: List[int] = []28        for val, (low, high), bins in zip(state, self.bounds, self.num_bins):29            clamped = max(low, min(high, val))30            frac = (clamped - low) / (high - low)31            bin_idx = min(bins - 1, int(frac * bins))32            bin_coords.append(bin_idx)33        return bin_coords34
35    def transform(self, state: List[float]) -> Dict[int, float]:36        """Map continuous state to a sparse 1-hot feature dictionary."""37        bin_coords = self.get_bin_indices(state)38        # Convert multi-dimensional coordinate to flat index39        flat_idx = 040        multiplier = 141        for coord, bins in reversed(list(zip(bin_coords, self.num_bins))):42            flat_idx += coord * multiplier43            multiplier *= bins44        return {flat_idx: 1.0}45
46    def evaluate(self, state: List[float], weights: List[float]) -> float:47        """Compute linear value estimate v_hat(s, w) = w^T x(s)."""48        sparse_x = self.transform(state)49        return sum(weights[idx] * val for idx, val in sparse_x.items())50
51
52def polynomial_feature_transform(state: List[float]) -> List[float]:53    """Expand a 2D continuous state [x1, x2] into quadratic polynomial features."""54    assert len(state) == 2, "Expected 2D input state."55    x1, x2 = state[0], state[1]56    return [1.0, x1, x2, x1 ** 2, x2 ** 2, x1 * x2]57
58
59# Verification matching the worked numerical example60transformer = StateAggregationTransformer(61    bounds=[(0.0, 3.0), (0.0, 3.0)],62    num_bins=[3, 3],63)64
65sample_state = [1.4, 2.3]66bin_coords = transformer.get_bin_indices(sample_state)67sparse_feat = transformer.transform(sample_state)68
69assert bin_coords == [1, 2]70assert sparse_feat == {5: 1.0}71
72grid_weights = [1.2, 2.5, 4.0, 2.0, 5.5, 8.0, 3.5, 7.0, 10.0]73grid_val = transformer.evaluate(sample_state, grid_weights)74assert abs(grid_val - 8.00) < 1e-475
76poly_feat = polynomial_feature_transform(sample_state)77assert poly_feat == [1.0, 1.4, 2.3, 1.96, 5.29, 3.22]78
79poly_weights = [0.5, 1.0, 1.5, 0.2, 0.1, 0.8]80poly_val = sum(f * w for f, w in zip(poly_feat, poly_weights))81assert abs(poly_val - 8.8470) < 1e-482
83print(f"State Aggregation Bin: {bin_coords} -> Flat Index: {list(sparse_feat.keys())[0]}")84print(f"Grid Discretization Value: {grid_val:.2f}")85print(f"Polynomial Features: {poly_feat}")86print(f"Polynomial Estimated Value: {poly_val:.4f}")87
88# -> State Aggregation Bin: [1, 2] -> Flat Index: 589# -> Grid Discretization Value: 8.0090# -> Polynomial Features: [1.0, 1.4, 2.3, 1.96, 5.29, 3.22]91# -> Polynomial Estimated Value: 8.8470

Watch Out For

The Quick Version

• Linear value approximators $\hat{v}(s, \mathbf{w}) = \mathbf{w}^\top \mathbf{x}(s)$ can only fit hyperplanes; feature construction embeds all non-linear domain complexity directly into $\mathbf{x}(s)$.
• Feature design controls the generalization-discrimination trade-off: broad features generalize quickly, while narrow features resolve sharp boundaries.
• Constructing interaction terms (e.g., $x_1 x_2$ or conjunctive receptive fields) is essential to capture diagonal synergies that purely additive models cannot express.
• Naive multi-dimensional grid discretization explodes exponentially ($B^D$), necessitating advanced sparse representations like tile coding or Fourier bases in higher dimensions.

What to Read Next

• Linear Methods in RL
• Coarse Coding in RL
• Tile Coding (CMAC) in RL
• Polynomials and Fourier Basis in RL