DQN save/restore not able to be restored

[shawnthesun123]

I am currently using the open_spiel DQN library (with colab, and Google Drive). I've trained a model as follows:

agent = dqn.DQN(
        session=sess,
        player_id=0,  # Assuming you're training for player 0
        state_representation_size=env.observation_spec()["info_state"][0],
        num_actions=num_actions,
        hidden_layers_sizes=hidden_layers_sizes,
        replay_buffer_capacity=replay_buffer_capacity,
        batch_size=batch_size)
    sess.run(tf.global_variables_initializer())

I then save the model using the save() function provided in the open_spiel dqn library as follows:
agent.save(checkpoint_dir)

However, when I try to restore the saved model using the same parameters and the restore() function provided in the open_spiel function using:
restored_agent.restore(checkpoint_dir)

I get this error: NotFoundError: Restoring from checkpoint failed. This is most likely due to a Variable name or other graph key that is missing from the checkpoint. Please ensure that you have not altered the graph expected based on the checkpoint.

After printing out all the variables in tf.trainable_variables(), I can see that the (to be) restored dqn agent always has 2 extra variables. Might be missing something, but I'm not quite sure what is going on

[lanctot]

Hi,

This is almost always because the TF environment (graph) looks different when restoring than when it was saved.

For TF to load properly, the graph has to be exactly the same. So e.g. if you trained with multiple copies of DQN, you have to recreate that environment with multiple DQN agents upon loading, even if you only intend to use one of them, because TF will recreate the graph the same way as during training.

Could this be it? If not, I will try locally tomorrow to see if there is an existing issue with saving/loading.

[shawnthesun123]

Thanks for the prompt response! I think I've only trained and uploaded one copy of DQN (I use model = dqn.DQN(), and then model.save()). I might have done it without knowing however - I can provide the current code I have if needed?

[lanctot]

Ok. Yes, please do. I will try to reproduce on my end, some time this week I hope.

[shawnthesun123]

Code to train and set DQN parameters:

checkpoint_dir = 'not/actual/link'

# Set DQN hyperparameters
hidden_layers_sizes = 64
replay_buffer_capacity = int(1e5)
batch_size = 32
num_train_episodes = int(1)

# Train the DQN model
with tf.Session() as sess:
    model = dqn.DQN(
        session=sess,
        player_id=0, 
        state_representation_size=env.observation_spec()["info_state"][0],
        num_actions=num_actions,
        hidden_layers_sizes=hidden_layers_sizes,
        replay_buffer_capacity=replay_buffer_capacity,
        batch_size=batch_size)
    sess.run(tf.global_variables_initializer())

    for ep in range(num_train_episodes):
        time_step = env.reset()
        while not time_step.last():
            action = model.step(time_step).action
            time_step = env.step([action, random.choice(range(num_actions))])  

        model.step(time_step)  
    
   model_variables = tf.trainable_variables()
   for var in model_variables:
      print(var)
   model.save(checkpoint_dir)

See variablesOnSave.txt
variablesOnSave.txt

Code to restore (placed in a separate Jupyter notebook cell then the above):

# Set DQN hyperparameters
hidden_layers_sizes = 64
replay_buffer_capacity = int(1e5)
batch_size = 32
num_train_episodes = int(1)

with tf.Session() as sess:
    restored_model = dqn.DQN(
        session=sess,
        player_id=0,
        state_representation_size=env.observation_spec()["info_state"][0],
        num_actions=num_actions,
        hidden_layers_sizes=hidden_layers_sizes,
        replay_buffer_capacity=replay_buffer_capacity,
        batch_size=batch_size)
    sess.run(tf.global_variables_initializer())
    model_variables = tf.trainable_variables()
    for var in model_variables:
      print(var)

    # Load the model from the checkpoint (ERROR OCCURS HERE)
    restored_model.restore(checkpoint_dir)

See variablesOnRestore.txt
variablesOnRestore.txt

Error:
fullerror.txt

[lanctot]

Hi Shawn, sorry for the delayed reply... I've been a bit tied up.

I ran the examples locally (not via colab). Specifically, these are the full files I ran:

• https://pastebin.com/p31QWa81 (train.py)
• https://pastebin.com/iPDPB75F (load.py)

And it works locally. train.py prints out:

<tf.Variable 'mlp/weights:0' shape=(196, 64) dtype=float32_ref>
<tf.Variable 'mlp/bias:0' shape=(64,) dtype=float32_ref>
<tf.Variable 'mlp/weights_1:0' shape=(64, 10) dtype=float32_ref>
<tf.Variable 'mlp/bias_1:0' shape=(10,) dtype=float32_ref>
<tf.Variable 'mlp_1/weights:0' shape=(196, 64) dtype=float32_ref>
<tf.Variable 'mlp_1/bias:0' shape=(64,) dtype=float32_ref>
<tf.Variable 'mlp_1/weights_1:0' shape=(64, 10) dtype=float32_ref>
<tf.Variable 'mlp_1/bias_1:0' shape=(10,) dtype=float32_ref>

and load.py prints out:

<tf.Variable 'mlp/weights:0' shape=(196, 64) dtype=float32_ref>
<tf.Variable 'mlp/bias:0' shape=(64,) dtype=float32_ref>
<tf.Variable 'mlp/weights_1:0' shape=(64, 10) dtype=float32_ref>
<tf.Variable 'mlp/bias_1:0' shape=(10,) dtype=float32_ref>
<tf.Variable 'mlp_1/weights:0' shape=(196, 64) dtype=float32_ref>
<tf.Variable 'mlp_1/bias:0' shape=(64,) dtype=float32_ref>
<tf.Variable 'mlp_1/weights_1:0' shape=(64, 10) dtype=float32_ref>
<tf.Variable 'mlp_1/bias_1:0' shape=(10,) dtype=float32_ref>

No errors and no additional variables.

Some ideas / follow-ups:

• Can you tell me what your version of TF is? I'm using 2.14.0.
• Can you run my simple example files based on yours above? Does it work? If so, maybe your original code is adding some variables unintentionally.
• Also... could it be that you're not including the disabling of TF2 line at the top? tf.disable_v2_behavior()

Hope this helps!

[shawnthesun123]

No worries!

I think something is going on with jupyter notebook - trying your example files in jupyter notebook resulted in the same error.

However, creating separate .py files resulted in no error, and worked perfectly!

I'm using version 2.14.0 as well. Here's my files and output:

Ipynb file:

file

Working:
train.py

load.py

Thank you so much for the help!