Skip to main content
2022-10-24 en

Engineering Reinforcement Learning for Edge Systems

Edge-system control creates two common reinforcement-learning bottlenecks: the observation can be a large collection of queues, capacities, and task descriptors, while the feasible action set can contain a huge number of discrete assignments. Compressing the observation and embedding the actions are attractive responses, but neither is automatically valid. The learned geometry must preserve what the controller needs to predict and choose.

This article treats state encoders, Wolpertinger-style action search, and variational autoencoders as engineering tools—not as results. Their value has to be established by controlled ablation against simpler baselines.

Representation quality means control sufficiency

Let an encoder map the available observation history to a latent state:

zt=eψ(ot,a<t).z_t=e_\psi(o_{\le t},a_{<t}).

A useful latent state need not reconstruct every byte of oto_t. It should retain enough information for reward prediction, transition prediction, and action selection. A practical approximation to this requirement is

p(rt,zt+1ot,at)p(rt,zt+1zt,at).p(r_t,z_{t+1}\mid o_{\le t},a_{\le t})\approx p(r_t,z_{t+1}\mid z_t,a_t).

If longer histories still improve next-step or return prediction after conditioning on ztz_t, the representation is probably aliased. That is a partial-observability problem, not proof that the physical system is unstable. The remedy may be recurrence, a belief-state model, or better timestamps and event features rather than a wider feed-forward network.

Three encoder objectives answer different questions:

ObjectivePreservesTypical blind spot
ReconstructionInformation needed to reproduce the observation under the chosen lossCan spend capacity on visually or numerically large but control-irrelevant detail
Reward and transition predictionFeatures useful for the modeled one-step dynamicsCan discard information needed for longer-horizon decisions or unseen actions
End-to-end controlFeatures that improve the current RL objectiveCan become unstable, task-specific, and hard to diagnose

A plain autoencoder minimizes a reconstruction risk such as

LAE(ψ,ω)=E[(ot,dω(eψ(ot)))].\mathcal{L}_{\mathrm{AE}}(\psi,\omega)=\mathbb{E}\left[\ell\left(o_t,d_\omega(e_\psi(o_t))\right)\right].

The loss \ell defines what “similar” means. Standardizing heterogeneous features, using masks for absent entities, and weighting units deliberately are therefore part of the model—not preprocessing trivia. A low reconstruction error does not establish that ztz_t is Markov or useful for control.

Predictive state learning instead trains ztz_t to explain reward and the next observation or next latent state. This can remove irrelevant detail, but it consumes transition data and inherits errors from the predictive model. The safest comparison is not “autoencoder versus predictive encoder” in isolation; it is each encoder with the same actor, critic, data budget, and evaluation seeds.

Entity structure should be explicit

Many edge observations are sets or sequences of similar entities. Flattening them into one vector makes parameter count depend on an arbitrary ordering and system size. Better options include:

  • a shared entity encoder followed by masked pooling or attention;
  • explicit entity identifiers only when roles are not exchangeable;
  • permutation tests that reorder equivalent entities without changing the decision;
  • separate global context from per-entity state instead of repeating it in every feature block.

Shared encoders improve statistical efficiency only when the entities actually share semantics. Different resource types, action constraints, or failure modes may require type embeddings or separate heads.

Encoder drift is a replay-buffer problem

If the encoder changes while the replay buffer stores latent vectors, an old ztz_t and a new ztz_t no longer represent observations in the same coordinate system. The critic then sees artificial non-stationarity created by the learning pipeline.

There are four defensible patterns:

  1. pretrain and freeze the encoder;
  2. store raw observations and re-encode them at sampling time;
  3. version latent features and invalidate incompatible replay;
  4. update the encoder slowly while measuring representation drift and refreshing targets.

Target encoders can smooth changes, but they do not make two coordinate systems identical. If a target encoder is used, log the distance between online and target embeddings and test whether critic targets change primarily because of the environment or because of the encoder.

Large discrete actions need meaningful geometry

The Wolpertinger architecture avoids evaluating every legal action. An actor first emits a continuous proto-action, an approximate nearest-neighbor index retrieves kk legal candidates, and the critic reranks them:

a^t=fθ(zt),Ck(a^t)=kNN(a^t,A),at=argmaxaCk(a^t)Qϕ(zt,a).\hat a_t=f_\theta(z_t),\qquad \mathcal{C}_k(\hat a_t)=\operatorname{kNN}(\hat a_t,\mathcal{A}),\qquad a_t^*=\arg\max_{a\in\mathcal{C}_k(\hat a_t)}Q_\phi(z_t,a).

This can make candidate retrieval sublinear in the number of actions, but only when the action embedding and nearest-neighbor index are useful. Euclidean proximity must correlate with similar consequences or value. An arbitrary binary code can make neighboring vectors represent unrelated assignments, leaving the true best action outside the candidate set.

The actor is optimized through a continuous surrogate while the environment executes a discrete candidate. That relaxation creates a training-execution mismatch. Validate it directly by measuring:

  • recall of the best action within the kk retrieved candidates on small instances where enumeration is possible;
  • regret between the selected candidate and the enumerated best critic action;
  • index-query time and critic reranking time separately;
  • sensitivity to kk, embedding dimension, and action-set changes;
  • the fraction of proto-actions near unsupported or infeasible regions.

Wolpertinger is only one design point. Factorized categorical policies, autoregressive assignment, hierarchical actions, action masking, and combinatorial optimization layers may better reflect the problem structure. Compare wall-clock latency and decision quality, not just the number of network outputs.

What a VAE actually optimizes

A variational autoencoder (VAE) introduces a probabilistic encoder qψ(zx)q_\psi(z\mid x), a prior p(z)p(z), and a decoder likelihood pω(xz)p_\omega(x\mid z). It maximizes an evidence lower bound:

LELBO=Eqψ(zx)[logpω(xz)]DKL(qψ(zx)p(z)).\mathcal{L}_{\mathrm{ELBO}}= \mathbb{E}_{q_\psi(z\mid x)}\left[\log p_\omega(x\mid z)\right] -D_{\mathrm{KL}}\left(q_\psi(z\mid x)\Vert p(z)\right).

The reconstruction term is not inherently mean-squared error. Its form follows the observation likelihood: a fixed-variance Gaussian yields a scaled squared-error term, while Bernoulli, categorical, count, or heteroscedastic likelihoods lead to different losses. Choosing a likelihood that matches normalized utilization, queue counts, categorical status, or censored measurements matters more than copying an image-VAE implementation.

The KL term regularizes the approximate posterior toward the prior and enables sampling. It does not guarantee a compact, disentangled, Markov, or control-sufficient latent state. A VAE is also not restricted to images, but its usefulness for structured system telemetry must be tested rather than inferred from image-generation success.

For control, compare a deterministic encoder and the VAE posterior mean before injecting latent sampling noise into the policy. Measure downstream return, critic calibration, transition prediction, and robustness under load shifts. Reconstruction quality alone is not a success criterion.

A disciplined experiment sequence

1. Establish simple baselines

Start with normalized raw features, a hand-designed summary, and a small multilayer perceptron. For actions, start with enumeration on reduced instances, factorized categorical outputs, or a domain heuristic. A complex latent model has to beat these baselines under the same interaction budget.

2. Test observability

Train probes for next reward and next observation from ot,ato_t,a_t, then repeat with short histories. If history produces a large gain, solve the observation problem before tuning the RL algorithm.

3. Separate representation from policy learning

Pretrain each encoder on the same logged transitions, freeze it, and train identical policies. Then run a second experiment with controlled fine-tuning. This distinguishes representation quality from optimization coupling.

4. Validate action retrieval on tractable cases

For small action spaces, enumerate all actions and calculate candidate recall and critic regret. Increase problem size only after the nearest-neighbor geometry passes this check.

5. Report ablations, not anecdotes

Vary latent dimension, candidate count kk, encoder update rate, and replay strategy one at a time. Use multiple seeds and report environment steps, wall-clock time, memory, return, and constraint or feasibility failures. Architecture sketches and isolated curves remain hypotheses until this comparison is complete.

Design conclusions

  • Compress observations only after defining which predictions and decisions the latent state must preserve.
  • Treat non-Markov observations as an information-design problem, not a synonym for unstable dynamics.
  • Make entity symmetry and action constraints explicit in the architecture.
  • Keep replay semantics consistent when encoders change.
  • Use Wolpertinger only when action-space geometry supports nearest-neighbor retrieval.
  • Choose a VAE likelihood from the data model; do not assume all VAEs use mean-squared error.
  • Present untested architecture proposals as hypotheses, never as experimental findings.

The strongest engineering result is often not a more elaborate network. It is a controlled demonstration that a representation preserves decision-relevant information and that an action-search approximation returns good feasible decisions within the system’s latency budget.

Further reading