在Tensorflow DQN中显示损失而不离开tf.Session()

import tensorflow as tf
import numpy as np

class DQN:
    def __init__(self, state_dim, action_dim, learning_rate):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.learning_rate = learning_rate
        
        self.inputs = tf.placeholder(tf.float32, [None, self.state_dim])
        self.targets = tf.placeholder(tf.float32, [None, self.action_dim])
        
        # 定义神经网络结构
        # ...
        
        self.loss = tf.reduce_mean(tf.square(self.targets - self.Q_values))
        self.optimizer = tf.train.AdamOptimizer(self.learning_rate).minimize(self.loss)

state_dim = 4  # 状态维度
action_dim = 2  # 动作维度
learning_rate = 0.001  # 学习率

dqn = DQN(state_dim, action_dim, learning_rate)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    
    for episode in range(num_episodes):
        # 获取当前状态state
        # ...
        
        # 根据当前状态选择动作action
        # ...
        
        # 执行动作，获取下一个状态next_state和奖励reward
        # ...
        
        # 计算目标Q值
        # ...
        
        # 更新神经网络参数
        _, loss = sess.run([dqn.optimizer, dqn.loss], feed_dict={dqn.inputs: state, dqn.targets: target_Q})
        
        # 显示损失
        print("Episode: {}, Loss: {}".format(episode, loss))