Source code for nnabla_rl.algorithms.mme_sac

# Copyright 2022,2023 Sony Group Corporation.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
from typing import Optional, Union

import gym

import nnabla_rl.model_trainers as MT
from nnabla_rl.algorithms import ICML2018SAC, ICML2018SACConfig
from nnabla_rl.algorithms.icml2018_sac import (DefaultExplorerBuilder, DefaultPolicyBuilder, DefaultQFunctionBuilder,
                                               DefaultReplayBufferBuilder, DefaultSolverBuilder,
from import ExplorerBuilder, ModelBuilder, ReplayBufferBuilder, SolverBuilder
from nnabla_rl.environments.environment_info import EnvironmentInfo
from nnabla_rl.models import QFunction, StochasticPolicy, VFunction
from nnabla_rl.utils.misc import sync_model

[docs]@dataclass class MMESACConfig(ICML2018SACConfig): """MMESACConfig List of configurations for MMESAC algorithm. Args: alpha_pi (Optional[float]): If None, will use reward_scalar to scale the reward. Otherwise 1/alpha_pi will be used to scale the reward. Defaults to None. alpha_q (float): Temperature value for negative entropy term. Defaults to 1.0. """ # override configurations reward_scalar: float = 5.0 alpha_pi: Optional[float] = None alpha_q: float = 1.0 def __post_init__(self): """__post_init__ Check the values are in valid range. """ super().__post_init__() if self.alpha_pi is not None: # Recompute with alpha_pi self.reward_scalar = 1 / self.alpha_pi
[docs]class MMESAC(ICML2018SAC): """Max-Min Entropy Soft Actor-Critic (MME-SAC) algorithm. This class implements the Max-Min Entropy Soft Actor Critic (MME-SAC) algorithm proposed by S. Han, et al. in the paper: "A Max-Min Entropy Framework for Reinforcement Learning" For details see: Args: env_or_env_info \ (gym.Env or :py:class:`EnvironmentInfo <nnabla_rl.environments.environment_info.EnvironmentInfo>`): the environment to train or environment info config (:py:class:`MMESACConfig <nnabla_rl.algorithms.mme_sac.MMESACConfig>`): configuration of the MMESAC algorithm v_function_builder (:py:class:`ModelBuilder[VFunction] <>`): builder of v function models v_solver_builder (:py:class:`SolverBuilder <>`): builder of v function solvers q_function_builder (:py:class:`ModelBuilder[QFunction] <>`): builder of q function models q_solver_builder (:py:class:`SolverBuilder <>`): builder of q function solvers policy_builder (:py:class:`ModelBuilder[StochasticPolicy] <>`): builder of actor models policy_solver_builder (:py:class:`SolverBuilder <>`): builder of policy solvers replay_buffer_builder (:py:class:`ReplayBufferBuilder <>`): builder of replay_buffer explorer_builder (:py:class:`ExplorerBuilder <>`): builder of environment explorer """ # type declarations to type check with mypy # NOTE: declared variables are instance variable and NOT class variable, unless it is marked with ClassVar # See for details _config: MMESACConfig def __init__(self, env_or_env_info: Union[gym.Env, EnvironmentInfo], config: MMESACConfig = MMESACConfig(), v_function_builder: ModelBuilder[VFunction] = DefaultVFunctionBuilder(), v_solver_builder: SolverBuilder = DefaultSolverBuilder(), q_function_builder: ModelBuilder[QFunction] = DefaultQFunctionBuilder(), q_solver_builder: SolverBuilder = DefaultSolverBuilder(), policy_builder: ModelBuilder[StochasticPolicy] = DefaultPolicyBuilder(), policy_solver_builder: SolverBuilder = DefaultSolverBuilder(), replay_buffer_builder: ReplayBufferBuilder = DefaultReplayBufferBuilder(), explorer_builder: ExplorerBuilder = DefaultExplorerBuilder()): super(MMESAC, self).__init__(env_or_env_info, config=config, v_function_builder=v_function_builder, v_solver_builder=v_solver_builder, q_function_builder=q_function_builder, q_solver_builder=q_solver_builder, policy_builder=policy_builder, policy_solver_builder=policy_solver_builder, replay_buffer_builder=replay_buffer_builder, explorer_builder=explorer_builder) def _setup_v_function_training(self, env_or_buffer): alpha_q = MT.policy_trainers.soft_policy_trainer.AdjustableTemperature( scope_name='alpha_q', initial_value=self._config.alpha_q) v_function_trainer_config = MT.v_value_trainers.MMEVTrainerConfig( reduction_method='mean', v_loss_scalar=0.5, unroll_steps=self._config.v_unroll_steps, burn_in_steps=self._config.v_burn_in_steps, reset_on_terminal=self._config.v_reset_rnn_on_terminal) v_function_trainer = MT.v_value_trainers.MMEVTrainer( train_functions=self._v, temperature=alpha_q, solvers={self._v.scope_name: self._v_solver}, target_functions=self._train_q_functions, # Set training q as target target_policy=self._pi, env_info=self._env_info, config=v_function_trainer_config) sync_model(self._v, self._target_v, 1.0) return v_function_trainer