RobertTLange / evosax

Evolution Strategies in JAX 🦎
Apache License 2.0
475 stars 44 forks source link

Strategy Save/Load Example #28

Open nhansendev opened 1 year ago

nhansendev commented 1 year ago

If I wanted to save/load the current strategy state and parameters so that I can resume optimization later, how would that be done?

Please let me know if there is existing documentation for this.

RobertTLange commented 1 year ago

Hi @Obliman! Yes, you can simply store/reload es_state and es_params and if the init args are correct, the strategy will continue from the previous break point. The init args don't need to match for all strategies, but for most of them (e.g. if recombination weights depends on the population size, etc.).

You can also simply pickle the strategy and reload the pickle object. I will leave this issue open and add an example/simple utility in the future. Does this help you? Cheers, Rob

nhansendev commented 1 year ago

Thanks for the response Rob. Yep, that's all I wanted to know.

nhansendev commented 1 year ago

You can also simply pickle the strategy and reload the pickle object.

Maybe I misunderstood, but as far as I can tell the states and params are not stored in the strategy instance and are just passed back and forth via arguments and returns, so would not be captured via directly pickling the instance.

To test it I tried pickling the strategy and state separately, resulting in a 384 B file for the strategy and a 36 kB file for the state, demonstrating that they remain separate.

To facilitate only needing to pickle the strategy I think the latest state and parameters could just be assigned as variables in the instance before saving e.g. self.saved_state=state, then retrieved after loading.

Does that sound about right?

RobertTLange commented 1 year ago

Heyo, that is correct. You would have to pickle both the strategy, es_state and es_params. I am thinking of adding a simple strategy wrapper that abstracts away some of the low-level control for the users that are not interested in mapping over different es_params, but just want a single strategy instance.

twoletters commented 1 year ago

Saving and restoring the state and parameters works perfectly to resume optimization. However, changing strategy or hyperparameters (e.g. population size or initial sigma) fails. Is it at all possible to do so by cherry-picking which parameters/states to reset and which to restore from the previous round of optimization?

twoletters commented 1 year ago

Is it at all possible to do so by cherry-picking which parameters/states to reset and which to restore from the previous round of optimization?

I was able to change sigma and the population size without too much trouble by updating the appropriate fields in the state, after restoring the latter from an earlier dump. However, it is strategy-dependent. I do not have a good solution to switch strategies at this point.

twoletters commented 1 year ago

Copying mean and optionally best_member, best_fitness and gen_counter from the old strategy's state to the new strategy's state is sufficient to hot-start the new strategy's optimization. Of course, this only hot-starts the mean and not the new strategy's internal states, such as memory or covariance matrix. I need to run tests to see if there is any merit to this method. Potentially, if a strategy is not performing well after a certain point, another could still further improve the results from there. IIRC, Hansen suggested to use SepCMA for 200 epochs before switching to VD-CMA to speed up the take-off.