Is there an algorithm that produces a uniform distribution over the set of trajectories with maximum reward sum?

Question

I am not an expert in reinforcement learning. I am applying it to my field of study.

I am training a model such that given a state, it predicts the probability of taking an action for every action possible. As such the model can be described as M: S -> P where P is the set of all probability vectors over A, the discrete action space. S is the set of all possible discrete states.

The model takes actions stochastically to create a finite trajectory (I do not choose the action with maximum probability, instead, I sample trajectories). Multiple trajectories have the same reward (summed over the state-action units in the trajectory).

Define G to be the set of trajectories with maximum reward sum.

If I train such a model using any mainstream algorithm (Q-learning, actor-critic, etc.) the model will learn, loosely speaking, to maximize the probability of a trajectory from G. It will not learn a uniform distribution over the set G. Instead, probability will concentrate on a specific trajectory. This is not an exploration problem as far as I know, since even if the model explores every possible trajectory, probability will still concentrate on one trajectory. And it is not inherently a problem in reinforcement learning, it is specific to my use case.

I would very much appreciate any references that tackle this problem. I am searching for an algorithm that produces a uniform distribution over the set G and a negligible probability for all other trajectories.

Answer 1

Q-learning will indeed learn a trajectory, as they are all equivalent, and usually the resulting policy of a Q-learning algorithm is a greedy policy, which selects: $$ a = argmax_a Q(s,a) $$ so, no matter what you do, the final policy in your case has to be deterministic.

In the actor critic scenario, this is not obvious. Indeed, if two actions $a_0$ and $a_1$ lead to trajectories with the same expected return, when you update the actor using: $$ \nabla J(\theta) \approx \mathbb{E}[G_t \nabla_\theta \ln \pi(a|s)] $$ you can clearly see that both the actions will be reinforced equally

However, you have to consider that the overall objective of RL is to maximise: $$ \mathbb{E}[\sum \gamma^tr_t] $$ In order to that to happen, going back to the previous scenario, any linear combination between the policy that always selects $a_0$ and the one that always selects $a_1$, will be an optimal policy, as the all lead to the same solution (solution meant as expected sum of rewards)

A very simple example of this, is by considering the trivial case where you have 2 bandits (think about slot machines) that give the same expected return $C$: your policy return will be $V = p(a_0)C + p(a_1)C$, clearly, any combination of the two probability lead to the same value (obviously having $p(a_0)C > 0, p(a_1) > 0, p(a_0) + p(a_1) = 1$)

However, luckily for you, there is an easy fix if you want the policy that not only maximizes the expected sum of rewards, but it's also as stochastic as possible, that is to introduce an entropy bonus $$ \nabla J(\theta) \approx \mathbb{E}[G_t \nabla_\theta \ln \pi(a|s)] + \mathcal{H}(\pi) $$ And this solution, will select the policy $\pi$ that also maximises the entropy, thus making it as stochastic as possible, thus distributing equally the probability mass across actions that lead to trajectories with the same expected return

Answer 2

I am not strictly convinced that Q-learning should necessarily learn a singular optimal path. Given infinite exploration, Q-learning should estimate the reward of all actions. In a state where two actions have the same reward the tie should be broken randomly (giving all equal paths the same probability of being followed). The remaining actions would have zero probability of being followed given a perfectly deterministic maximizing agent.

In practice we do not explore infinitely or our estimators do not learn infinitely. This might account for the perception that it finds a singular near optimal path. Often the problem space is too large and our estimators too opaque to inspect what has actually been learned. I would challenge to experiment with tabular Q-learning on a multi-armed bandit where more than one arm has equivalent rewards.

So I guess my answer would be that the issue is unlikely to be with the learning method but with the learning hyper-parameters (learning rate, epsilon annealing rate, etc).