can one specify parts of the model that are non differentiable? about evojax HOT 2 CLOSED

# Only ask the solver for the parameters that belong to the non-diff operations.
solver = NeuroEvolutionAlgo(param_size=non_diff_param_size)

# Define the policy to contain both diff and non-diff operations.
class AwesomePolicy:

   def __init__(self):
       # policy network definition and init
       # Save the part that are differentiable.
       self.diff_params = xxxxx

  def reset():
      # Save the differentiable params in the policy state.
      return PolicyState(params=self.diff_params)

   def get_action(self, params, p_state):
       # The "params" parameter is coming from the solver, they belong to the non-diff ops.
        load_diff_params_to_model(p_state.params)
        load_non_diff_params(params)
        # Model inference