stochastic policy used for exploration outperforms its deterministic version at test time

tldr: how can i make an agent learn a stochastic policy that can preserve its performance even if converted into a deterministic one?

i’m training a policy with PPO in a task with continuous action space, during training in each timestep the agent uses a neural network to compute the mean of a gaussian distribution then samples the action from it while during test/evaluation phase it just uses the mean computed deterministically as action, afaik this is the standard approach with on-policy algos

however, i noticed that in my experiment the return computed using the data in the rollout buffer (hence using the experience collected by the stochastic and explorative version of policy) was sensibly higher than the return computed running the same policy but with no sampling, which is supposed to be entirely focused on exploitation and therefore perform better

i’ve also thought that it could be a bug in return computation so i tested again the policy but sampling actions instead of using the mean, just like in training, i also rendered the episode to check what was happening and well, in this case the performance was indeed drastically superior, the agent could recover from bad states and handle difficult situations much better than its deterministic version

my guess is that during training the stochastic nature of the policy became a fundamental part of it and the agent was relying on that, if my guess is correct the learned policy was inherently stochastic and using only the mean was like using another different distribution, too different for a on-policy method

“aight then just keep sampling actions” some might say, but deploying what is basically a explorative policy doesn’t really sound a good idea to me lol, especially in environments where you may need fast and precise actions like in robotics

yea in some environments i could still sample actions and sanify them to be sure they won’t break anything if executed but now i’m wondering if this kind of performance gap you get while converting a gaussian policy to its deterministic version has been studied before and if there are “conversion aware” methods that would learn a stochastic policy but with the guarantee that the learned policy would still perform at least as good as the original gaussian one once converted in a deterministic policy and also without premature convergence to local optima

i got a couple ideas to try, like changing the sign of the entropy coefficient after some steps, or using something like a epsilon-greedy action selection mechanism to alternate between sampling and using the mean or also introducing a scheduler for the variance to push the policy towards a delta distribution, but id like to hear more ideas