RobertTLange / evosax

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

best member is not updated when using centered_rank=True #32

Closed ynotzort closed 1 year ago

ynotzort commented 1 year ago

Hey!

When optimizing with a FitnessShaper with centered_rank=True, then the state of the optimizer does not update best_member as well as best_fitness since at each generation the best_fitness == -0.5. The issue arises because of the '<' in this line here: https://github.com/RobertTLange/evosax/blob/db3f1aab11a4c87c131010116d2353406c8d4afd/evosax/utils/helpers.py#L15

replacing it with '<=' would replace the best_member at each generation. However I am not sure if this is actually better than the current behavior...

The current behavior makes this setting not usable with some strategies such as GLD...

RobertTLange commented 1 year ago

Hi @ynotzort! Yes, that is true and a known problem. I am not entirely sure how this can be circumvented. One option would be to specify the fitness shaping arguments together with the strategy at instantiation and then to always do the fitness shaping internally within the strategy.tell call. So something along these lines:

fitness_config = {"centered_rank"=True,
                           "z_score"=False,
                           "w_decay"= 0.0}
strategy = CMA_ES(popsize=16, num_dims=2, fitness_config=fitness_config)
...
es_state = strategy.tell(x, fitness, es_state, es_params)

This way the strategy will always have access to the raw fitness. The original fitness shaping outside of the strategy can still be applied, but in that case the default fitness_config should not apply any additional shaping inside the strategy. What do you think and would you interested in opening a PR?

ynotzort commented 1 year ago

Hey!

Yes, that could be done, but then all the strategies would probably need to implement something like

tell_strategy(
        self,
        x: chex.Array,
        fitness: chex.Array,
        state: EvoState,
        params: EvoParams,
        fitness_raw: Optional[chex.Array] = None,  # <-- new
    )

But that would just push the problem onto the individual strategy implementations...

RobertTLange commented 1 year ago

Fixed in PR #34