Yes, it is possible to have an RL agent output logits instead of the final action probabilities. To achieve this, you need to adapt both the action space of the environment and the actor network to work with logits rather than probabilities. You can achieve this by modifying the actor network and the environment's action space definition as mentioned below:
Modifications to the Actor Network:
To have the actor network output logits instead of probabilities, ensure that the network does not include a softmax layer at the end. This alteration will allow the network to output logits directly.
Modification to the Environment’s Action Definition:
Since the actor network will output logits, the environment's action space definition should be adjusted accordingly. You can define the action space using an "rlNumericSpec" object with dimensions corresponding to the number of actions (the number of logits) output by the actor network. For more details on the "rlNumericSpec" object, refer to: https://www.mathworks.com/help/reinforcement-learning/ref/rl.util.rlnumericspec.html.
By implementing these changes, the PPO agent will output logits instead of probabilities. You can handle the masking and apply the softmax function within your custom environment. Here is an example of how you can set up the PPO agent:
actorNet = create_actor_network(numObs, numActions); % No softmax layer at the end
criticNet = create_critic_network(numObservations); % Create critic network
agentOpts = rlPPOAgentOptions('ActorNetwork', actorNet, 'CriticNetwork', criticNet);
agent = rlPPOAgent(observationInfo, actionInfo, agentOpts);
This setup should help resolve your issue. If you provide more specific details about your setup or include some of your code, I can offer more targeted advice.
