# # Unity ML-Agents Toolkit
# ## ML-Agent Learning (SAC)
# Contains an implementation of SAC as described in https://arxiv.org/abs/1801.01290
# and implemented in https://github.com/hill-a/stable-baselines
import numpy as np
from mlagents.trainers.tf_policy import TFPolicy
from mlagents.trainers.brain import BrainParameters
LOGGER = logging . getLogger ( " mlagents.trainers " )
logger = logging . getLogger ( " mlagents.trainers " )
BUFFER_TRUNCATE_PERCENT = 0.8
"""
def __init__ (
self , brain , reward_buff_cap , trainer_parameters , training , load , seed , run_id
self ,
brain_name : str ,
reward_buff_cap : int ,
trainer_parameters : dict ,
training : bool ,
load : bool ,
seed : int ,
run_id : str ,
: param brain_name : The name of the brain associated with trainer config
: param reward_buff_cap : Max reward history to track in the reward buffer
: param trainer_parameters : The parameters for the trainer ( dictionary ) .
: param training : Whether the trainer is set for training .
: param load : Whether the model should be loaded .
super ( ) . __init__ ( brain , trainer_parameters , training , run_id , reward_buff_cap )
super ( ) . __init__ (
brain_name , trainer_parameters , training , run_id , reward_buff_cap
)
self . param_keys = [
" batch_size " ,
" buffer_size " ,
]
self . check_param_keys ( )
self . load = load
self . seed = seed
self . policy : TFPolicy = None
self . step = 0
self . train_interval = (
if " save_replay_buffer " in trainer_parameters
else False
)
self . sac_policy = SACPolicy (
seed , brain , trainer_parameters , self . is_training , load
)
self . policy = self . sac_policy
# Load the replay buffer if load
if load and self . checkpoint_replay_buffer :
try :
self . load_replay_buffer ( )
except ( AttributeError , FileNotFoundError ) :
LOGGER . warning (
" Replay buffer was unable to load, starting from scratch. "
)
LOGGER . debug (
" Loaded update buffer with {} sequences " . format (
self . update_buffer . num_experiences
)
)
for _reward_signal in self . policy . reward_signals . keys ( ) :
self . collected_rewards [ _reward_signal ] = defaultdict ( lambda : 0 )
def save_model ( self ) - > None :
def save_model ( self , name_behavior_id : str ) - > None :
"""
Saves the model . Overrides the default save_model since we want to save
the replay buffer as well .
"""
Save the training buffer ' s update buffer to a pickle file.
"""
filename = os . path . join ( self . policy . model_path , " last_replay_buffer.hdf5 " )
LOGGER . info ( " Saving Experience Replay Buffer to {} " . format ( filename ) )
filename = os . path . join (
self . trainer_parameters [ " model_path " ] , " last_replay_buffer.hdf5 "
)
logger . info ( " Saving Experience Replay Buffer to {} " . format ( filename ) )
with open ( filename , " wb " ) as file_object :
self . update_buffer . save_to_file ( file_object )
"""
filename = os . path . join ( self . policy . model_path , " last_replay_buffer.hdf5 " )
LOGGER . info ( " Loading Experience Replay Buffer from {} " . format ( filename ) )
filename = os . path . join (
self . trainer_parameters [ " model_path " ] , " last_replay_buffer.hdf5 "
)
logger . info ( " Loading Experience Replay Buffer from {} " . format ( filename ) )
LOGGER . info (
logger . info (
" Experience replay buffer has {} experiences. " . format (
self . update_buffer . num_experiences
)
)
if trajectory . done_reached :
self . _update_end_episode_stats ( agent_id )
self . _update_end_episode_stats (
agent_id , self . get_policy ( trajectory . behavior_id )
)
def is_ready_update ( self ) - > bool :
"""
self . update_sac_policy ( )
self . update_reward_signals ( )
def create_policy ( self , brain_parameters : BrainParameters ) - > TFPolicy :
policy = SACPolicy (
self . seed ,
brain_parameters ,
self . trainer_parameters ,
self . is_training ,
self . load ,
)
for _reward_signal in policy . reward_signals . keys ( ) :
self . collected_rewards [ _reward_signal ] = defaultdict ( lambda : 0 )
# Load the replay buffer if load
if self . load and self . checkpoint_replay_buffer :
try :
self . load_replay_buffer ( )
except ( AttributeError , FileNotFoundError ) :
logger . warning (
" Replay buffer was unable to load, starting from scratch. "
)
logger . debug (
" Loaded update buffer with {} sequences " . format (
self . update_buffer . num_experiences
)
)
return policy
def update_sac_policy ( self ) - > None :
"""
Uses demonstration_buffer to update the policy .
num_updates = self . trainer_parameters [ " num_update " ]
batch_update_stats : Dict [ str , list ] = defaultdict ( list )
for _ in range ( num_updates ) :
LOGGER . debug ( " Updating SAC policy at step {} " . format ( self . step ) )
logger . debug ( " Updating SAC policy at step {} " . format ( self . step ) )
buffer = self . update_buffer
if (
self . update_buffer . num_experiences
for stat , stat_list in batch_update_stats . items ( ) :
self . stats_reporter . add_stat ( stat , np . mean ( stat_list ) )
bc_module = self . sac_ policy. bc_module
bc_module = self . policy . bc_module
if bc_module :
update_stats = bc_module . update ( )
for stat , val in update_stats . items ( ) :
# Get minibatches for reward signal update if needed
reward_signal_minibatches = { }
for name , signal in self . policy . reward_signals . items ( ) :
LOGGER . debug ( " Updating {} at step {} " . format ( name , self . step ) )
logger . debug ( " Updating {} at step {} " . format ( name , self . step ) )
# Some signals don't need a minibatch to be sampled - so we don't!
if signal . update_dict :
reward_signal_minibatches [ name ] = buffer . sample_mini_batch (
update_stats = self . sac_policy . update_reward_signals (
update_stats = self . policy . update_reward_signals (
reward_signal_minibatches , n_sequences
)
for stat_name , value in update_stats . items ( ) :
def add_policy ( self , name_behavior_id : str , policy : TFPolicy ) - > None :
"""
Adds policy to trainer .
: param brain_parameters : specifications for policy construction
"""
if self . policy :
logger . warning (
" add_policy has been called twice. {} is not a multi-agent trainer " . format (
self . __class__ . __name__
)
)
self . policy = policy
def get_policy ( self , name_behavior_id : str ) - > TFPolicy :
"""
Gets policy from trainer associated with name_behavior_id
: param name_behavior_id : full identifier of policy
"""
return self . policy
GitHub
4 年前