Reinforcement Learning method suitable for a large discrete action space with high sample efficiency

Question

Consider the following problem.

We have a process, that generates $N$ stones (e.g. 2000) in one batch $b$. Every pebble has state $s_{i}^b$ and reward $s_i^b$. After choosing one pebble $i$ from the $N$, we start sampling again using the chosen pebble as a starting point and we generate the next batch $b+1$. The state $s_i^b$ is a vector of real-values and the reward $r_i^b$ is a real value.

The problem is to choose pebbles so that we maximize reward $r_i^b$ in long term. Because depending on how we choose the pebble, we can sample around the region that gives a better or worse reward $r$.

During each new batch, we make one selection of one pebble (so actions can from $i, \dots, N$. We have access to the previous $m$ batches (e.g. through replay buffer) with their rewards and states.

In short, it looks like this:

0. Chose randomly the first pebble from which we start sample;
1. We start sampling from the chosen pebble in the current batch;
2. We sample $N$ pebbles from a process, each pebble have a state $s_i^b$ and reward $r_i^b$;
3. We can chose one pebble from $i \dots N$ as action $a_i^b$ based on state $s_i^b$ and reward $r_i^b$;
4. Go to point 1 and repeat;

For example, at the moment, I choose in a given batch $b$ pebble with max reward $r_i^b$, so

$$i = \underset{i}{\mathrm{argmax}}\, r_i^b$$ and then use $a_i^b$ for a the current batch $b$.

But what I want is to choose: $$i = \underset{i}{\mathrm{argmax}}\, \underset{b}{E}[R_i^b | s_i, a_i]$$

Graphically speaking:

Assuming one Batch (N) is 30 P - pebble, P - chosen pebble

batch 1: PPPPPPPPPPPPPPPPPPPPPPPPPPPPPP

batch 2: PPPPPPPPPPPPPPPPPPPPPPPPPPPPPP

batch 3: PPPPPPPPPPPPPPPPPPPPPPPPPPPPPP

batch 4: PPPPPPPPPPPPPPPPPPPPPPPPPPPPPP

batch 5: PPPPPPPPPPPPPPPPPPPPPPPPPPPPPP

So, if I have a batch $N$, when I choose one element from the batch as an action, so the expected reward is the highest in long term, and start the sampling again from it. So only one element per batch can be chosen. And only choice of the one pebble per batch does affect the sequence in the next batch, but not inside the current batch.

The problem is, what Reinforcement Learning algorithm to use, when we choose only one item from N. And the one choice affect the whole sampled sequence in the next batch. For example in batch 1 to 4, the reward can be very low, and in batch 5 the reward is super high, if we chose wisely the pebbles in 4 previous batchs.

Answer 1

My understanding of your environment is:

• The batch number $b$ is the same thing as a time step $t$. Each batch is associated with a single static representation of the environment, the agent makes one decision, then receives a reward and a next state. I am using $t$ instead of $b$ to match the usual name seen in RL.

• There is a generator which can be in a state $s_t$, initially unknown for $s_0$.

• Each time step, $t$, the generator creates a batch of "pebbles". The generator state $s_t$ determines the output of the generator, and is the only non-random influence on the contents of each batch - it controls the possible rewards plus the possible next states.

• Selecting a specific pebble from the batch returns an immediate reward $r_{t+1}$ plus sets the next state $s_{t+1}$. This is the agent's action, and it can choose any pebble from the batch as its action $a_t$.

• The agent can see the values of $r_{t+1}$ and $s_{t+1}$ associated with each pebble before making its selection.

This is a slightly unusual setup for an MDP, but I think it is still a valid MDP, and the expected return can still be maximised using reinforcement learning. One important detail is that although you have a discrete list of actions in each batch, the action index values look like they are meaningless. Your action choice is really to choose the next state $s_{t+1}$ from a presented list of available states. Your action space is therefore continuous, not discrete, even though at each step you are selecting from a discrete set of actions.

The environment design is different enough from standard that most off-the-shelf RL agent libraries will not work for you. However, you can take advantage of this knowledge of the action space, and the fact that there is no other meaningful state than that selected by the agent when it chooses a pebble, to learn state values efficiently.

I would suggest you use a variant of Q-Learning, based on the "afterstate" values that you are selecting. You can train a neural network to approximate the state value function $\hat{v}(s,\theta)$ and you have a choice whether to use discounted returns or an average reward setting. When deciding the greedy target policy, you will want to select according to

$$\pi = \text{argmax}_i (r_i + \gamma \hat{v}(s_i,\theta))$$

for discounted return. For Q learning, technically you should also store the whole batch output in experience replay, because the TD target will be

$$g = \text{max}_i (r_i + \gamma \hat{v}(s_i,\theta))$$

and you won't know what this max will be later at the time of storing the memory. There may be ways around that - for instance you could use SARSA updates based on the actual action taken, or you could store the value of $g$. However, these would be less sample efficient and you would need to discard the memory at a faster rate, before it became too out of date with respect tothe current policy.

I think an average reward will be closer to your stated goal of generating highest expected reward from each individual batch, but I am not 100% certain how to set up the policy for that. Whilst a discounted return may end up preferring a high immediate reward in some circumstances. However, you may be able to get near identical behaviour by setting $\gamma$ high enough (e.g. $0.99$ or higher) in the more familiar discounted return setting.

The efficiency of this approach will depend on how easy it is for the agent to learn the mapping between states and the quality of the next batch. However, you should have a higher sample efficiency when selecting by afterstate, because the agent does not need to learn separately how actions relate to states.