PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Action Representations for Reinforcement Learning

Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, Philip S. Thomas

arXiv:1902.00183v2cs.LGstat.ML

TL;DR

Reinforcement learning often lacks learned structure over large action sets, limiting generalization when many discrete choices are available. The paper factorizes the policy into an internal action-representation policy and an embedding-to-action mapping, learning representations from observed transitions with supervised updates. The resulting framework provides optimal-policy representation and convergence guarantees and is demonstrated on large-scale real-world problems.

  • Problem

    Reinforcement learning methods either ignore action-space structure or rely on predefined representations, despite real-world problems involving large finite action sets.

  • Method

    The policy is factorized into an internal policy over a low-dimensional action-representation space and a function mapping representations to actual actions, with the mapping learned from observed transitions.

  • Results

    The overall-policy class can represent an optimal policy, and the combined supervised-learning and policy-gradient procedure has almost-sure asymptotic convergence guarantees.

  • Takeaways & Limitations

    Learning action representations from observed transitions supports generalization over large action sets and can be combined with existing policy-gradient methods while preserving convergence guarantees.

  • Takeaways & Limitations

    The main results focus on finite action sets, and the existence theorem does not ensure that all optimal policies are representable by the overall-policy class.

Abstract

from arXiv · show

Most model-free reinforcement learning methods leverage state representations (embeddings) for generalization, but either ignore structure in the space of actions or assume the structure is provided a priori. We show how a policy can be decomposed into a component that acts in a low-dimensional space of action representations and a component that transforms these representations into actual actions. These representations improve generalization over large, finite action sets by allowing the agent to infer the outcomes of actions similar to actions already taken. We provide an algorithm to both learn and use action representations and provide conditions for its convergence. The efficacy of the proposed method is demonstrated on large-scale real-world problems.

1. Introduction

The paper targets reinforcement-learning problems with large action sets by learning action representations rather than treating actions as unrelated choices. It introduces a factorized policy, learning procedure, convergence analysis, and real-world evaluations.

  • Large action sets limit reinforcement learning in real-world decisions such as trading, power regulation, and personalized tutoring.
  • The overall policy combines an internal policy acting in action-representation space with a function mapping representations to actual actions.
  • The method learns action representations from observed transitions, unlike prior approaches that assume fixed, predefined representations.
  • Action-representation learning is framed as supervised learning, allowing updates to generalize feedback across actions with similar impacts on the state.The supervised component learns the mapping function while the internal policy is trained with policy-gradient methods.
  • The overall-policy parameterization can represent an optimal policy within the proposed policy class.
  • The paper proves equivalence of policy-gradient updates, establishes almost-sure asymptotic convergence, and evaluates the approach on recommendation domains with thousands of actions.

2. Background

The paper formulates reinforcement learning as a finite-action Markov decision process and defines policies, rewards, transitions, and value functions within that framework.

  • The environment is modeled as a discrete-time Markov decision process with discrete states and finite actions.The notation uses M = (S, A, P, R, γ, d0).
  • The action set A is finite, while the main results can extend to continuous state spaces.The paper denotes the action-set size by |A| and restricts its focus to finite action sets.
  • P defines transition probabilities, R gives expected rewards for state-action pairs, γ discounts future rewards, and d0 specifies the initial-state distribution.
  • A policy is a conditional distribution over actions given the current state.The paper writes this distribution as π(a|s).
  • The objective is to find a policy maximizing the expected sum of discounted future rewards, characterized by state-action and state value functions.

3. Related Work

Related work reduces or exploits action structure through factorization, auxiliary supervision, motor primitives, or output representations. This paper instead learns action representations within an RL policy using supervised updates.

  • Factorizing Action Space: Action-space factorization methods reduce large action spaces using binary or category-based decompositions, but often require handcrafted decompositions.
  • Factorizing Action Space: Continuous-action policy gradients can select the nearest discrete action, and later work extended action-representation lookup to larger domains.
  • Auxiliary Tasks: Auxiliary-task research uses supervised transition prediction to learn state representations or intrinsic rewards.
  • Auxiliary Tasks: The proposed approach similarly uses a supervised loss, but applies it to an action-representation module within the overall policy.
  • Motor Primitives: Motor-primitive research decomposes behavior into mid-level abstractions, while related policy compositions use local policy-gradient information.
  • Other Domains: Unlike output embeddings in supervised learning, this work treats policy outputs as actions and learns representations that capture useful action correlations for reinforcement learning.

4. Generalization over Actions

The paper factorizes policy decisions through a learned, low-dimensional action-representation space, enabling generalization across actions while preserving optimal-policy expressiveness under stated assumptions. PG-RA learns the action mapping from transition data with supervised learning and updates the internal policy using policy gradients, with convergence established through multi-timescale analysis.

  • Generalization over Actions: The factorized policy samples an action representation with πi and deterministically maps it to an action with f.The representation space is E ⊆ R^d, while f maps E to the action set A.
  • Existence of an Optimal Policy: Under Assumptions (A1)–(A2), there exist πi and f such that the overall policy has optimal state value vπo = v⋆.The assumptions include deterministic action decoding from embeddings and sufficient transition information to infer the executed action.
  • Supervised Learning of f: The action mapping f is learned by estimating the executed action from transition tuples and minimizing a supervised loss based on KL divergence.The estimators reconstruct P(At|St, St+1), using observed transitions rather than reward signals.
  • Policy-Gradient Learning of πi: The internal policy can be trained directly with its policy gradient, avoiding explicit computation of the potentially intractable inverse mapping f −1.Lemma 2 establishes that the internal-policy update still optimizes the overall policy’s performance.
  • PG-RA Algorithm and Convergence: PG-RA initializes action representations from random-policy trajectories using the supervised loss, optionally keeps them fixed, and can learn both components simultaneously.The simultaneous-learning procedure is supported by a three-timescale convergence analysis.

5. Empirical Analysis

The empirical analysis evaluates action representations in a maze and two real-world recommender MDPs with large action sets. Learned representations preserve action structure and improve policy performance and learning speed over standard actor-critic methods.

  • 5.1. Domains: The maze uses a continuous state, an exponentially large action set from binary actuators, a long horizon, and a sparse goal reward.Each actuator is independently on or off, and the net displacement is the sum of selected actuator displacements; noise affects actions 10% of the time.
  • 5.2. Visualizing the Learned Action Representations: The learned maze embeddings preserve relative action-transition structure: nearby representations correspond to similar displacement magnitude and direction.Smooth color transitions reflect preserved structure; Tanh bounds the representation coordinates to [−1, 1], producing edge squashing.
  • 5.3. Performance Improvement: Standard actor-critic performance deteriorates as maze action counts increase, whereas action-representation variants perform consistently across settings.Figure 5 compares maze domains with 2^4, 2^8, and 2^12 actions; shaded regions show one standard deviation across 10 trials.
  • 5.3. Performance Improvement: The proposed variants achieve up to 2× and 3× higher return on the tutorial and software MDPs, respectively, while learning faster.Standard actor-critic methods tend to select one-step actions with high returns, while the proposed methods handle longer time horizons.
  • 5.3. Performance Improvement: Supervised updates learn action representations, leaving only internal-policy parameters to high-variance policy-gradient updates.The paper attributes faster learning to reduced update variance and reports that the internal policy can more quickly approximate an optimal policy.

6. Conclusion

The paper concludes that exploiting structure in action spaces enhances generalization for large action sets in real-world applications. It emphasizes simplicity, theoretical representational guarantees, and extensibility across policy-gradient methods.

  • Simplicity: Observed transitions support an additional supervised update rule that learns action representations.The conclusion presents this as the approach’s simplicity advantage.
  • Theory: The overall policy class can represent an optimal policy, and the paper derives learning procedures for its parameters.This is the stated theoretical advantage of the approach.
  • Extensibility: PG-RA can be extended with other policy-gradient methods while preserving convergence guarantees.The conclusion identifies this extensibility as a key advantage.

A. Proof of Lemma 1

The proof of Lemma 1 introduces a representation variable and uses the policy’s conditional-independence structure to show that an overall policy can match the original policy’s value.

  • A. Proof of Lemma 1: Lemma 1 states that under Assumptions (A1)–(A2), every policy has an internal policy πi compatible with the function f.The resulting construction underlies the overall-policy representation.
  • A. Proof of Lemma 1: A representation variable e is introduced through the law of total probability to rewrite the policy’s performance.The derivation inserts e and then factors terms using conditional distributions.
  • A. Proof of Lemma 1: The derivation uses conditional independence of the next state from the representation given the state and action.This is the transition assumption used to simplify the expanded expression.
  • A. Proof of Lemma 1: The Markov property makes the action conditionally independent of the state given the representation.This supports replacing the original action distribution with the representation-conditioned formulation.

B. Proof of Lemma 2

The proof of Lemma 2 rewrites the overall-policy gradient through the internal representation policy and establishes the corresponding expected parameter updates. It then places the convergence analysis within a multi-timescale framework.

  • B. Proof of Lemma 2: Lemma 2 considers deterministic functions f that map representation points e ∈ R^d to actions and analyzes expected updates to θ.The lemma formulates the update using the internal policy while the mapping to actions remains deterministic.
  • B. Proof of Lemma 2: Using Lemma 1, the proof expresses the overall-policy performance in terms of the internal policy and differentiates it with respect to θ.The derivation applies the policy-gradient theorem to the resulting parameterization.
  • B. Proof of Lemma 2: Because e deterministically maps to a, the action-value function under the overall policy equals the internal-policy action value at e.This identity permits the gradient expression to be rewritten over the representation space.
  • B. Proof of Lemma 2: The PG-RA convergence proof adapts existing multi-timescale results to three dependent parameter sequences.The analysis first reviews actor-critic convergence and then establishes asymptotic behavior across three timescales.

C.2. Three-Timescale Setup

The analysis uses three timescales so X and Y track their stationary points while Z evolves most slowly. Under Assumptions B1–B5, the coupled iterates converge to the corresponding stationary configuration.

  • Three-timescale updates: The algorithm extends two-timescale actor-critic analysis with a three-timescale approach for updating the action prediction module.The three-timescale procedure is chosen to make convergence faster than an inner-loop supervised-update alternative.
  • Stationary points: Z converges to Z⋆ when X and Y provide converged estimates between successive Z updates.Assumption B5 supplies the stationary-point condition for the slowest recursion.
  • Three-timescale updates: X updates on the fastest timescale, Y on an intermediate timescale, and Z on the slowest timescale.This ordering makes slower variables quasi-static while faster recursions approach their stationary points.
  • Stationary points: With fixed Z, Y converges to µ1(Z), and with fixed Y and Z, X converges to µ2(Y, Z).These stationary-point properties follow from Assumptions B3 and B4.
  • Convergence: Under Assumptions B1–B5, (X_t, Y_t, Z_t) → (µ2(µ1(Z⋆), Z⋆), µ1(Z⋆), Z⋆) as t →∞.The proof treats the recursions as noisy Euler discretizations and shows asymptotic errors vanish.

C.3. PG-RA Convergence Using Three- Timescales:

The PG-RA updates are mapped to a three-timescale stochastic approximation system whose assumptions support convergence. The internal policy parameters converge almost surely under the stated assumptions.

  • Assumptions: The projection operator keeps the internal policy parameters bounded, satisfying the boundedness requirement for the convergence analysis.The other parameter sets are assumed bounded as well.
  • Convergence: Under Assumptions A1–A5, the internal policy parameters θ_t converge to Ẑ as t →∞ with probability one.The three-timescale mapping allows the result from Lemma 3 to establish this theorem.
  • Update structure: PG-RA jointly updates the critic, action representation modules, and internal policy through three dependent stochastic recursions.The TD error depends on critic, action-representation, and internal-policy parameters because rewards, bootstrapping, and the baseline couple the updates.
  • Assumptions: For fixed internal-policy parameters, the action representation updates follow supervised learning, while a linear critic converges to its TD stationary point for fixed policy representations.The critic argument uses TD(λ), with λ = 0 in the algorithm and the proof extending to λ > 0.
  • Assumptions: The B5 condition can be relaxed from a global to a local asymptotically stable fixed point because convergence is the required conclusion.With optimal critic and action representations, the internal policy follows the overall policy gradient and converges to its local fixed point.

D. Implementation Details

The implementation estimates action probabilities stochastically because computing the action-prediction expectation requires integrating over the representation space. It parameterizes action similarity with learnable embeddings and a temperature-controlled Boltzmann distribution.

  • Action prediction: Computing the action prediction module requires a complete integral over e, so the implementation uses a stochastic estimate without a closed-form solution.Sample-based evaluation can become computationally expensive as the dimension of e increases.
  • Action prediction: The estimator ˆf assigns action probabilities from similarity between each action representation and the embedding produced by ˆg(s, s′).W contains learnable action representations as columns, while z_a measures similarity to the state-transition embedding.
  • Action prediction: A Boltzmann distribution with temperature τ converts similarity scores into valid action probabilities.As τ →0, the distribution concentrates on the action with the most similar representation, yielding the deterministic estimate for ˆf.

D.2. Hyper-parameters

The experiments use different network, discounting, representation-dimension, and initialization settings for the maze and real-world environments. Baselines receive additional depth and width searches to account for PG-RA’s deeper architecture.

  • Maze domain: In the maze domain, actor and critic networks are single-layer, action representations have dimension 2, and 2000 random trajectories initialize the representation.The discount factor is γ = 0.99 and λ = 0.9.
  • Real-world environments: In real-world environments, actor and critic networks are two-layer, hidden widths are searched over {64, 128, 256}, and 10,000 random trajectories initialize action representations.Representation dimensions are searched over {16, 32, 64}, and the action predictor is improved online.
  • Real-world environments: Real-world experiments use γ = 0.9, λ = 0.9 for actor-critic methods, and a DPG target update rate of 0.001.The internal policy is parameterized as an isotropic normal distribution with searched variance settings.
  • Baseline tuning: PG-RA’s policy decomposition resembles a one-layer-deeper neural network, so baselines are searched over additional depths {1, 2, 3}.Baseline widths are searched over {2, 16, 64}, with remaining search properties matched to PG-RA.
Loading 1902.00183v2…