class Optimizer:
    def __init__(self, env, policy_network, model_path, params_path='params.pkl'):
        self.env = env
        self.policy_network = policy_network
        self.model_path = model_path
        self.params_path = params_path

    def objective(self, trial, n_episodes=100):
        lr = trial.suggest_loguniform('lr', 1e-5, 1e-1)
        gamma = trial.suggest_uniform('gamma', 0.9, 0.999)
        optimizer = optim.Adam(self.policy_network.parameters(), lr=lr)
        trainer = REINFORCE(self.env, self.policy_network, optimizer, self.model_path, gamma=gamma)
        reward = trainer.train(n_episodes, save_model=False, save_video=False)
        return reward

    def optimize(self, n_trials=100, save_params=True):
        if not TRAIN and os.path.isfile(self.params_path):
            with open(self.params_path, 'rb') as f:
                best_params = pickle.load(f)
            print("Loaded parameters from disk")
        elif not FINETUNE:
            best_params = {'lr': LEARNING_RATE, 'gamma': GAMMA}
            print(f"Using default parameters: {best_params}")
        else:
            print("Optimizing hyperparameters")
            study = optuna.create_study(direction='maximize')
            study.optimize(self.objective, n_trials=n_trials)
            best_params = study.best_params
            if save_params:
                with open(self.params_path, 'wb') as f:
                    pickle.dump(best_params, f)
                print("Saved parameters to disk")
        return best_params