Fix imports for reinforcement learning