diff --git a/gym_stirap/envs/stirap_env.py b/gym_stirap/envs/stirap_env.py
index fccd9a14f0c7fbc017963a27d454657cb29a5011..1be42b4a186858c56d4b28e603bc713726fb7dc1 100644
--- a/gym_stirap/envs/stirap_env.py
+++ b/gym_stirap/envs/stirap_env.py
@@ -262,6 +262,8 @@ class StirapEnv(gym.Env):
 
             self.probability_line, = plt.plot(self.x, ys)
             self.potential_line, = plt.plot(self.x, 0.5*self.dx * self.potential(),'k')
+            plt.fill_between(self.x, 2 * np.abs(self.target_psi)**2 * self.dx, color='C2',alpha=.3)
+            
             plt.ylim([0, .2])
 
             if timeevolution: