How to access policy model in step function of environment in rllib

Hello,
I am using DQN algorithm and I need to obtain the (optimal) action related to a specific state in the step function of environment. Is there any way to do it?