Maximum Entropy Softmax Policy Gradient via Entropy Advantage Estimation

1 Introduction

Entropy regularisation is pivotal to many practical deep reinforcement learning (RL) algorithms. Practical algorithms such as Trust Region Policy Optimization (TRPO) [Schulman et al., 2015a] penalise the policy improvement or greedy step using Kullback-Leibler (KL) divergence (also called as relative entropy) to regularise the deviations between consecutive policies. This method, often termed KL regularisation, has been the foundational approach for contemporary deep RL algorithms [Vieillard et al., 2020; Geist et al., 2019].

Another critical approach, Maximum Entropy RL (MaxEnt RL) [Ziebart, 2010; Haarnoja et al., 2018; Levine, 2018; Marino et al., 2021; Cetin and Celiktutan, 2022], augments the conventional RL task objective with an entropy term to directing policies toward areas of higher expected trajectory entropy. It is known to improve the exploration and robustness of policies by promoting stochasticity [Eysenbach and Levine, 2019; Eysenbach and Levine, 2021]. In practice, this can be implemented by adding an entropy reward to the original task reward.

1 Technical appendix and codes available at: https://github.com/ rnilva/eapo- ijcai25

Recent theoretical advancements have shown the effectiveness of MaxEnt RL in accelerating the convergence of policy gradient (PG) methods [Mei et al., 2020; Agarwal et al., 2021; Cen et al., 2022]. However, despite the enticing theoretical support, a significant gap exists between theory and practice. While off-policy methods like Soft Actor-Critic [Haarnoja et al., 2018] have demonstrated practical success with MaxEnt RL, its application in simpler direct Softmax policy gradient settings remains surprisingly underexplored.

We hypothesise that this research gap is potentially attributed to the difficulty of handling the entropy reward in practice. [Yu et al., 2022] empirically analysed the problematic nature of the entropy reward using SAC. Authors pointed out that in an episodic setting, the entropy return is largely correlated to the episode's length, thereby rendering the policy overly optimistic or pessimistic, and even in infinitehorizon settings, the entropy reward can still obscure the task reward. This raises a fundamental question: can MaxEnt RL provide practical benefits in straightforward stochastic Softmax policy gradient settings?

Inspired by this observation, we proposed a simple but practical approach to control the impact of the entropy reward. In this paper, we introduce Entropy Advantage Policy Optimisation (EAPO), a method that estimates the task and entropy objectives of the regularised (soft) objective separately. By employing a dedicated discount factor for the entropy reward and utilising Generalised Advantage Estimation (GAE) [Schulman et al., 2015b] on each objective separately, EAPO controls the effective horizon of the entropy return estimation and the entropy regularisation on policy evaluation. EAPO's simplicity requires only minor modifications to existing advantage actor-critic algorithms, providing a clear demonstration of impact of MaxEnt RL framework in basic settings. Our empirical evaluation not only extend the wellestablished PPO [Schulman et al., 2017b] and TRPO [Schulman et al., 2015a], but also demonstrate superior stability in discrete action spaces compared to TD-SAC [Zhou et al., 2024], a discrete version of SAC with improved stability.

Figure 1 illustrates the challenge of learning the MaxEnt policy for an episodic task using a naive implementation that simply augments the task reward with an entropy reward. In this task, the agent is required to reach the goal state while performing the minimum number of actions. The naive MaxEnt agent fails to learn the optimal stochastic policy, resulting in two failure modes: acting almost deterministically when the temperature τ is low or wandering around indefinitely when τ is high. In contrast, EAPO successfully achieves the near-optimal stochastic policy by utilising TD(0) learning [Sutton and Barto, 2018] (i.e., set GAE λ to 0) for the entropy objective. Additionally, the example demonstrates that lowering the discount factor for the entropy estimation γH helps prevent the inflation of the entropy reward [Yu et al., 2022] and reduces sensitivity to the temperature.

In this work, we primarily focus on empirical examinations of the Softmax MaxEnt Policy Gradient method [Levine, 2018; Mei et al., 2020] across diverse environments: 4 discretised [Tang and Agrawal, 2020] MuJoCo continuous control tasks, 16 Procgen episodic environments [Cobbe et al., 2020], and the MiniGrid DoorKey environment [ChevalierBoisvert et al., 2023]. Our results demonstrate the efficacy of MaxEnt RL policies in improving optimisation performance, generalization, and potential exploration benefits.

2 Background

2.1 Preliminaries

This work considers a finite undiscounted Markov Decision Process (MDP) ⟨S,A, r, ρ, T ⟩, where S is the set of states s and A is the set of actions a, and ρ is the initial state distribution. T is the transition function T : S ×A ↦→∆(S), where ∆(S) is the probability simplex over S, and r is the reward function r : S × A ↦→ R. We introduce discount factors γV and γH as variance reduction parameters as in [Schulman et al., 2015b]. We define the value function of state s under the policy π as V π (s) := Es0=s,at∼π(·|st),st+1∼T (st,at) [∑∞t=0 γ tV r(st, at)]. Also, the action-value of performing action a at state s under the policy π is Qπ (s, a) := Es0=s,a0=a,at>0∼π(·|st),st+1∼T (st,at) [∑∞t=0 γ tV r(st, at)]. And define the advantage function Aπ as Aπ (s, a) := Qπ (s, a) - Eπ(·|s) [Qπ (s, ·)] = Qπ (s, a) - V π (s). We also define the cumulative discounted entropy return of state s under policy π as

In deterministic MDPs, the cumulative discounted entropy return V π H (s) represents the Shannon entropy of the possible future trajectories' distribution [Levine, 2018; Tiapkin et al., 2023], and we refer to it as trajectory entropy throughout this work for brevity.

The objective of Maximum Entropy Reinforcement Learning (MaxEnt RL), or often Regularised MDPs [Geist et al., 2019; Neu et al., 2017] is to maximise the expectation of the sum of the value and the trajectory entropy with respect to the initial state distribution:

where the temperature parameter τ ≥ 0 is a hyperparameter to be controlled to balance the significance between these two objectives, and we introduce the distinct discount factors. We ensure the objective J(π) is finite for all policies by assuming either an episodic MDP or an MDP with an absorbing state that yields zero reward.

2.2 Soft Advantage Function

Analogous to the definition of the action-value function Qπ as the expected cumulative rewards after selecting an action a [Sutton and Barto, 2018], we define QπH as the expected future trajectory entropy after selecting an action:

The definition arises naturally from the consideration that uncertainty exists due to the stochastic policy at the current state, which has settled by the time an action is performed. Consequently, the QπH is simply the expected discounted future trajectory entropy.

From the recursive relation of trajectory entropy from (1) and the defintion (4), the following relation is derived:

We now define the entropy advantage function AπH analogous to the conventional advantage function:

We let ˜ V π (s) := V π (s) + τV π H (s) as the soft value function, and let ˜Qπ (s, a) := Qπ (s) + τQπH(s, a) as the soft Qfunction. Finally, we define the soft advantage function:

2.3 Soft Policy Gradient Theorem

[Shi et al., 2019] showed that it is possible to optimise the soft objective using direct policy gradient from samples. Thus, we can use the soft advantage function to find the policy that maximises the MaxEnt RL objective.

Theorem 1 (Soft Policy Gradient). Let J(π) the MaxEnt RL objective defined in 2. And πθ(a|s) be a parameterised policy. Then,

We provide the proof in the appendix. While the exact soft policy gradient theorem requires the corresponding exponential discount term for each advantage estimate, we use the approximate policy gradient in this work by replacing ˜ Aπγ (st, at) with ˜ Aπ (st, at). It is worth noting that when the exact gradient is known, [Mei et al., 2020] proved that the soft policy gradient has the global convergence property and may converge faster than the policy gradient without entropy regularisation despite the objective being biased. However, in our practical setup, this is not guaranteed.

3 RelatedWorks

3.1 MaxEnt RL in Discrete Action Spaces

Off-policy methods like Soft Q-Learning and SAC [Haarnoja et al., 2017; Schulman et al., 2017a] learn soft Q-functions for discrete action spaces but often suffer from instability due to the coupled challenges of temperature adjustment, off-policy soft Q-function estimation, and policy updates [Christodoulou, 2019; Zhou et al., 2024; Xu et al., 2021]. In contrast, EAPO utilises direct policy gradient and approximates the soft value function ˜ V π (s) using GAE from onpolicy samples, making it less prone to the estimation error at the cost of sample efficiency. This work also demonstrates the better policy optimisation stability and performance of EAPO over SD-SAC [Zhou et al., 2024] across procgen environments.

3.2 Soft Policy Gradient Method

A more directly related work is [Shi et al., 2019], which explored a soft policy gradient method emphasising its inherent simplicity and proved the soft policy gradient theorem. Their method also involves estimating ˜Qπ from off-policy samples and introduces additional techniques for mitigating the estimation issues. On the other hand, EAPO uses the soft advantage function ˜ Aπ for policy gradient estimator with additional hyperparameters to reduce variance and seamlessly integrates with existing techniques, such as value function normalisation, due to its structural equivalence between its method for estimating the entropy advantage function and the conventional advantage function.

3.3 Reward Inflation Problem

[Yu et al., 2022] showed that the entropy reward of MaxEnt RL can cause reward inflation from indefinite exploration in episodic settings as it is given at every time step before termination. We show that having a lower discount factor for entropy objective can mitigate the problem and establish positive implications for the use of entropy rewards in policy optimisation.

3.4 Entropy Cost vs. MaxEnt RL

A more common approach to applying entropy regularisation to PG methods is to add an entropy cost term to the sample-based policy gradient estimator to maximise the policy entropy at each sampled state, retaining the stochasticity of the policy during optimisation process [Mnih et al., 2016; Schulman et al., 2017b]. While this entropy bonus term remains a heuristic approach despite its practical success [Ahmed et al., 2019], MaxEnt RL provides a theoretically grounded framework that directs a policy toward regions of higher expected trajectory entropy, albeit at the cost of bias imposed on the objective [Levine, 2018; Schulman et al., 2017a]. This work demonstrates the performance improvements of the MaxEnt approach of EAPO over PPO's entropy bonus term.

4 ProposedMethod

4.1 Overview

In this section, we develop our Entropy Advantage Policy Optimisation (EAPO) method. At its core, EAPO independently estimates both the value advantage function and the entropy advantage function and combines them to derive the soft advantage function. EAPO adopts a separate prediction head to the conventional value critic to approximate the trajectory entropy of a state, which is then used for entropy advantage estimation. We extend the PPO [Schulman et al., 2017b] and TRPO [Schulman et al., 2015a] by substituting the advantage estimate with the soft advantage estimate and omitting the entropy bonus term.

4.2 Entropy Advantage Estimation

The entropy advantage AπH is estimated from the sampled log probabilities of the behaviour policy. We utilise the Generalised Advantage Estimation (GAE) [Schulman et al., 2015b]

for a variance-reduced estimation of the entropy advantage:

where δH t := -logπ(at|st)+γHV π H (st+1)-V π H (st), and γH and λH are the discount factor and GAE lambda for entropy advantage estimation, respectively. Note that the equation is the same as the GAE for the conventional advantage, except the reward term is replaced by the negative log probability. This simplicity is also consistent with the remark that the only modification required for the MaxEnt policy gradient is to add the negative log probability term to the reward at each time step [Levine, 2018].

4.3 Entropy Critic

An entropy critic network, parameterised by ω, approximates the trajectory entropy V π H (·; ω). To train the network, we construct the TD(λH) bootstrapped target ˆ V π H (st; ω) from sampled trajectories. We estimate the entropy advantages using the entropy critic, and the target is derived as ˆ V π H (st; ω) = ˆ AH,GAE(λH,γH) (st, at; ω) + V π H (st; ω). Then the entropy critic is trained by minimising the mean squared error using the semi-gradient method: LH(ω) := ˆEt[ 1 2 (V π H (st; ω) - ˆ V π H (st; ω′ ))2], where ω′ indicates it is

treated as a constant. Note that this process is exactly the same procedure of training the standard critic, except we use negative log probabilities instead of rewards.

Throughout the conducted experiments, we implemented the entropy critic network to share its parameters with the return value critic V V ϕ , with only the final linear layers for outputting its prediction distinct. This form of parameter sharing allows minimal computational overhead to implement EAPO.

Further, we employ the PopArt normalisation [van Hasselt et al., 2016] to address the scale difference of entropy and return estimates. It is important to note that the negative log probability -logπ(at|st) is collected for every timestep. In contrast, the reward can be sparse, leading to significant magnitude variations based on the dynamics of the environment [Hessel et al., 2019]. This discrepancy can pose challenges, especially when using a shared architecture. Thus, utilising the value normalisation technique like PopArt is pivotal for the practical implementation of EAPO.

4.4 Entropy Advantage Policy Optimisation

Subsequently, we integrate the entropy advantage with the standard advantage estimate ˆ AπV , also computed using GAE and return value critic parameterised by ϕ, analogously to the entropy advantage estimation process we describe above. Then the soft advantage function ˜ Aπ is

where ˆ AV,GAE(λV ,γV ) is the value advantage estimation using GAE. Finally, we substitute the estimated conventional advantage function in the policy objective of PPO and TRPO with ˜ Aπ . The PPO objective function becomes:

where rθ t is the probability ratio between the behaviour policy πθold (at|st) and the current policy πθ(at|st), and c1, c2 and ϵ are hyperparameters to be adjusted. The value critic loss LV is also defined by the mean square error, LV (ϕ) = ˆEt[ 12 (V (st; ϕ) - ˆ V π (st))2] where ˆ V π is the return value

estimate.

Similarly, the optimisation problem of TPRO becomes:

where δ is a hyperparameter.

4.5 Combining KL and Entropy Regularisation

EAPO, by extending TRPO, combines KL regulasization on policy updates and entropy regularisation on policy evaluation. Recent works [Vieillard et al., 2020; Geist et al., 2019; Shani et al., 2020] suggest such combination could lead to improved convergence in regularised MDPs. Exploring the practical implications of this combined regularisation remains an interesting direction for future work.

5 Experiments

In this section, we evaluate the policy optimisation performance of EAPO against the corresponding baseline on-policy algorithms, PPO and TRPO. Specifically, we assess the optimisation efficiency for episodic tasks and the generalisation capability of EAPO on 16 Procgen [Cobbe et al., 2020] benchmark tasks. Moreover, we investigate EAPO's efficacy on continuing control tasks using 4 discretised popular MuJoCo [Todorov et al., 2012] tasks, and we analyse the impact of hyperparameters τ, γH and λH. Finally, we include MiniGrid-DoorKey-8x8 task [Chevalier-Boisvert et al., 2023] to examine if EAPO can help solve the hard exploration task.

We implemented EAPO using Stable-baselines3 [Raffin et al., 2019] and conducted experiments on environments provided by the Envpool [Weng et al., 2022] library. All empirical results are averaged over 10 random seeds, with 95% confidence intervals indicated.

For the hyperparameter selection, we conducted a brief search for baseline algorithm hyperparameters that perform reasonably well, tuning only the EAPO-specific hyperparameters such as γH to ensure fair comparisons. Implementation details and hyperparameters are reported in the appendix.

5.1 Procgen Benchmark Environments

We evaluate PPO-based EAPO on the 16 Procgen benchmark environments [Cobbe et al., 2020], which feature discrete actions and image-based observations. These environments include tasks with varied correlations between the episode

Mean Normalised Score

EAPO τ:0.02

EAPO τ:0.005

PPO c:0.01

PPO c:0.001

Train (EAPO)

Test (EAPO)

Train (EAPO)

Test (EAPO)

Train (PPO)

Test (PPO)

Train (PPO)

Test (PPO)

0.0

0.5

1.0

1.5

2.0

2.5

Timesteps

1e7

Figure 2: Left: Procgen test results of EAPO with γH = 0.8, λH = 0.95, and τ ∈ (0.02, 0.005) against PPO with entropy coefficients c ∈ (0.001, 0.01). Right: The mean normalised score for both test and training.

length and the return, making them suitable for testing MaxEnt RL algorithms. Following [Cobbe et al., 2020], we train on 200 procedurally generated levels and test on 100 unseen levels using easy difficulty setting. The same set of hyperparameters is used across all environments. We also compare against SD-SAC [Zhou et al., 2024], a discrete version of SAC [Haarnoja et al., 2018] that adopts a clipping mechanism for Q-function update and introduces entropy penalty to mitigate policy entropy collapsing. For hyperparameter tuning, we selected three representative environments: BigFish (long survival), Climber (indefinite exploration) and Dodgeball (survival followed by quick goal-reaching). These tasks represent different episode length-return correlations.

Figure 2 and Table 1 summarise the generalisation test results of EAPO, baseline PPO agents with varying τ and entropy bonus coefficients, and SD-SAC. The mean normalised score is computed as in [Cobbe et al., 2020]. EAPO with γH = 0.8 and λH = 0.95 significantly surpasses the baseline PPO, achieving a 29% higher mean normalised score (0.54 ± 0.06 vs 0.42 ± 0.07) and consistently outperforming across most environments, with particularly large improvements in Plunder (+116%), BigFish (+55%), and Dodgeball (+60%). These results demonstrate EAPO's effectiveness in both test (generalisation) and training (optimisation) phases. Notably, SD-SAC significantly underperforms in these environments, highlighting the difficulty of stabilising discrete SAC. SDSAC particularly struggles with the reward inflation problem, as showen by the its pattern of performance: it achieves reasonable results only in environments with positive lengthreturn correlation (StarPilot, Fruitbot and BigFish) where extended episodes naturally lead to higher returns, while failing in environments allowing indefinite actions without reward (e.g., Plunder, Leaper). We provide the used hyperparameters and the full learning curves of SD-SAC in the appendix.

Moreover, we investigated the impact of the GAE λH for the entropy advantage estimation, finding it does not significantly affect performance. This suggests adjusting γH and τ

is usually sufficient in episodic tasks given a small λH.

Figure 2 (Right) shows the improved generalisation of high-entropy policies. Higher temperature τ favours highentropy trajectories (see Figure 3), performing similarly or worse during the training but better in testing. This aligns with [Eysenbach and Levine, 2021], showing MaxEnt policies' robustness to distributional shifts.

Figure 4 demonstrates that the lower discount factor γH mitigates the reward inflation problem [Yu et al., 2022]. A small γH significantly improves the performance in environments where agents can traverse without meaningful reward gain (Dodgeball and Climber) while maintaining performance in environments inherently requiring longer episodes (Bigfish, Bossfight).

5.2 Discretised Continuous Control Tasks

We measure the performance of EAPO extending PPO (EAPO-PPO) and TRPO (EAPO-TRPO) on continuing control tasks in 4 MuJoCo environments, comparing them against their corresponding baselines. For the PPO baselines, we searched for the best entropy coefficient within the set c ∈ (0.0001, 0.001, 0.01). Additionally, we tested the PPO and TRPO agent with the reward augmented by the entropy reward -τ logπ(at|st) to evaluate the impact of separating the MaxEnt objective. Note that the entropy rewardaugmented baseline is effectively regarded as EAPO with γH = γ and λH = λ, but without the entropy critic. We discretise the continuous action space using the method proposed by [Tang and Agrawal, 2020]. Results using the original continuous action space are provided in the appendix. We measured the mean episodic return of the stochastic policy periodically over 100 episodes during the training. We compare EAPO to the PPO agent with the best-performing entropy coefficient, and with the entropy reward augmented PPO.

The training curves are presented in Figure 5. The result shows that by adjusting γH and λH, we can configure EAPO

Mean Normalisied Score

0.8

0.7

0.6

0.5

0.4

0.3

0.2

0.1

0.0

to outperform or match the conventional entropy regularisation method throughout all environments. We found that the best-performing values of γH and λH vary depending on the characteristics of the environment, similar to their value counterparts γ and λ, respectively. Although EAPO demonstrates more stable performance compared to the entropy bonus, this relatively modest performance gain suggests that EAPO may be less efficient for continuing tasks. Figure 5 also demonstrates that the adjustability adopted by EAPO improves the naive implementation of the MaxEnt policy that augments the entropy reward. We also provide ablation experiments on γH and λH using PPO-based EAPO in the appendix.

5.3 MiniGrid-DoorKey-8x8 Environment

Finally, we evaluate the exploration performance of PPObased EAPO on the MiniGrid-DoorKey-8x8 environment [Chevalier-Boisvert et al., 2023]. Figure 6 shows that EAPO achieves consistent success, solving this hard exploration task within 5M frames across all 10 seeds, while the baseline PPO succeeds in only 3 seeds. However, EAPO did not improve the stability since both EAPO and PPO were highly sensitive to their common hyperparameters (such as γV and λV ). Moreover, EAPO's improvement in success rate appears to contradict to the theory that entropy regularisation may not effectively mitigate epistemic uncertainty [Mei et al., 2020].

One possible explanation is that EAPO's trajectory entropy estimation, based on log probabilities from on-policy samples, tends to underestimate the entropy of previously visited trajectories. This underestimation induces an implicit bias toward exploring new trajectories.

6 Conclusion

We have introduced EAPO, a model-free on-policy actorcritic algorithm based on the maximum entropy reinforcement learning framework. Our approach facilitates a practical MaxEnt RL algorithm by taking advantage of existing mechanisms for standard value learning in actor-critic algorithms to the entropy objective. Through empirical evaluations, EAPO has been shown to replace the entropy cost method and that a more principled entropy maximisation method enhances policy optimisation. Moreover, by providing a stable MaxEnt method for discrete action spaces, EAPO enables deeper investigation of the MaxEnt RL in emerging applications such as GFlowNets [Mohammadpour et al., 2024; Tiapkin et al., 2024] and RLHF [Christiano et al., 2017]. We anticipate that the method's simplicity will encourage broader adoption of MaxEnt RL and inference-based methods [Cetin and Celiktutan, 2022; Marino et al., 2021; Ward et al., 2019] in promising areas like competitive RL and robust RL.