stable/freqtrade/freqai/RL/ReinforcementLearnerCustomAgent.py

263 lines
9.7 KiB
Python

# import logging
# from pathlib import Path
# from typing import Any, Dict, List, Optional, Tuple, Type, Union
# import gym
# import torch as th
# from stable_baselines3 import DQN
# from stable_baselines3.common.buffers import ReplayBuffer
# from stable_baselines3.common.policies import BasePolicy
# from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor
# from stable_baselines3.common.type_aliases import GymEnv, Schedule
# from stable_baselines3.dqn.policies import CnnPolicy, DQNPolicy, MlpPolicy, QNetwork
# from torch import nn
# from freqtrade.freqai.data_kitchen import FreqaiDataKitchen
# from freqtrade.freqai.RL.BaseReinforcementLearningModel import BaseReinforcementLearningModel
# logger = logging.getLogger(__name__)
# class ReinforcementLearnerCustomAgent(BaseReinforcementLearningModel):
# """
# User can customize agent by defining the class and using it directly.
# Here the example is "TDQN"
# Warning!
# This is an advanced example of how a user may create and use a highly
# customized model class (which can inherit from existing classes,
# similar to how the example below inherits from DQN).
# This file is for example purposes only, and should not be run.
# """
# def fit_rl(self, data_dictionary: Dict[str, Any], dk: FreqaiDataKitchen):
# train_df = data_dictionary["train_features"]
# total_timesteps = self.freqai_info["rl_config"]["train_cycles"] * len(train_df)
# policy_kwargs = dict(activation_fn=th.nn.ReLU,
# net_arch=[256, 256, 128])
# # TDQN is a custom agent defined below
# model = TDQN(self.policy_type, self.train_env,
# tensorboard_log=str(Path(dk.data_path / "tensorboard")),
# policy_kwargs=policy_kwargs,
# **self.freqai_info['model_training_parameters']
# )
# model.learn(
# total_timesteps=int(total_timesteps),
# callback=self.eval_callback
# )
# if Path(dk.data_path / "best_model.zip").is_file():
# logger.info('Callback found a best model.')
# best_model = self.MODELCLASS.load(dk.data_path / "best_model")
# return best_model
# logger.info('Couldnt find best model, using final model instead.')
# return model
# # User creates their custom agent and networks as shown below
# def create_mlp_(
# input_dim: int,
# output_dim: int,
# net_arch: List[int],
# activation_fn: Type[nn.Module] = nn.ReLU,
# squash_output: bool = False,
# ) -> List[nn.Module]:
# dropout = 0.2
# if len(net_arch) > 0:
# number_of_neural = net_arch[0]
# modules = [
# nn.Linear(input_dim, number_of_neural),
# nn.BatchNorm1d(number_of_neural),
# nn.LeakyReLU(),
# nn.Dropout(dropout),
# nn.Linear(number_of_neural, number_of_neural),
# nn.BatchNorm1d(number_of_neural),
# nn.LeakyReLU(),
# nn.Dropout(dropout),
# nn.Linear(number_of_neural, number_of_neural),
# nn.BatchNorm1d(number_of_neural),
# nn.LeakyReLU(),
# nn.Dropout(dropout),
# nn.Linear(number_of_neural, number_of_neural),
# nn.BatchNorm1d(number_of_neural),
# nn.LeakyReLU(),
# nn.Dropout(dropout),
# nn.Linear(number_of_neural, output_dim)
# ]
# return modules
# class TDQNetwork(QNetwork):
# def __init__(self,
# observation_space: gym.spaces.Space,
# action_space: gym.spaces.Space,
# features_extractor: nn.Module,
# features_dim: int,
# net_arch: Optional[List[int]] = None,
# activation_fn: Type[nn.Module] = nn.ReLU,
# normalize_images: bool = True
# ):
# super().__init__(
# observation_space=observation_space,
# action_space=action_space,
# features_extractor=features_extractor,
# features_dim=features_dim,
# net_arch=net_arch,
# activation_fn=activation_fn,
# normalize_images=normalize_images
# )
# action_dim = self.action_space.n
# q_net = create_mlp_(self.features_dim, action_dim, self.net_arch, self.activation_fn)
# self.q_net = nn.Sequential(*q_net).apply(self.init_weights)
# def init_weights(self, m):
# if type(m) == nn.Linear:
# th.nn.init.kaiming_uniform_(m.weight)
# class TDQNPolicy(DQNPolicy):
# def __init__(
# self,
# observation_space: gym.spaces.Space,
# action_space: gym.spaces.Space,
# lr_schedule: Schedule,
# net_arch: Optional[List[int]] = None,
# activation_fn: Type[nn.Module] = nn.ReLU,
# features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
# features_extractor_kwargs: Optional[Dict[str, Any]] = None,
# normalize_images: bool = True,
# optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
# optimizer_kwargs: Optional[Dict[str, Any]] = None,
# ):
# super().__init__(
# observation_space=observation_space,
# action_space=action_space,
# lr_schedule=lr_schedule,
# net_arch=net_arch,
# activation_fn=activation_fn,
# features_extractor_class=features_extractor_class,
# features_extractor_kwargs=features_extractor_kwargs,
# normalize_images=normalize_images,
# optimizer_class=optimizer_class,
# optimizer_kwargs=optimizer_kwargs
# )
# @staticmethod
# def init_weights(module: nn.Module, gain: float = 1) -> None:
# """
# Orthogonal initialization (used in PPO and A2C)
# """
# if isinstance(module, (nn.Linear, nn.Conv2d)):
# nn.init.kaiming_uniform_(module.weight)
# if module.bias is not None:
# module.bias.data.fill_(0.0)
# def make_q_net(self) -> TDQNetwork:
# # Make sure we always have separate networks for features extractors etc
# net_args = self._update_features_extractor(self.net_args, features_extractor=None)
# return TDQNetwork(**net_args).to(self.device)
# class TMultiInputPolicy(TDQNPolicy):
# def __init__(
# self,
# observation_space: gym.spaces.Space,
# action_space: gym.spaces.Space,
# lr_schedule: Schedule,
# net_arch: Optional[List[int]] = None,
# activation_fn: Type[nn.Module] = nn.ReLU,
# features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
# features_extractor_kwargs: Optional[Dict[str, Any]] = None,
# normalize_images: bool = True,
# optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
# optimizer_kwargs: Optional[Dict[str, Any]] = None,
# ):
# super().__init__(
# observation_space,
# action_space,
# lr_schedule,
# net_arch,
# activation_fn,
# features_extractor_class,
# features_extractor_kwargs,
# normalize_images,
# optimizer_class,
# optimizer_kwargs,
# )
# class TDQN(DQN):
# policy_aliases: Dict[str, Type[BasePolicy]] = {
# "MlpPolicy": MlpPolicy,
# "CnnPolicy": CnnPolicy,
# "TMultiInputPolicy": TMultiInputPolicy,
# }
# def __init__(
# self,
# policy: Union[str, Type[TDQNPolicy]],
# env: Union[GymEnv, str],
# learning_rate: Union[float, Schedule] = 1e-4,
# buffer_size: int = 1000000, # 1e6
# learning_starts: int = 50000,
# batch_size: int = 32,
# tau: float = 1.0,
# gamma: float = 0.99,
# train_freq: Union[int, Tuple[int, str]] = 4,
# gradient_steps: int = 1,
# replay_buffer_class: Optional[ReplayBuffer] = None,
# replay_buffer_kwargs: Optional[Dict[str, Any]] = None,
# optimize_memory_usage: bool = False,
# target_update_interval: int = 10000,
# exploration_fraction: float = 0.1,
# exploration_initial_eps: float = 1.0,
# exploration_final_eps: float = 0.05,
# max_grad_norm: float = 10,
# tensorboard_log: Optional[str] = None,
# create_eval_env: bool = False,
# policy_kwargs: Optional[Dict[str, Any]] = None,
# verbose: int = 1,
# seed: Optional[int] = None,
# device: Union[th.device, str] = "auto",
# _init_setup_model: bool = True,
# ):
# super().__init__(
# policy=policy,
# env=env,
# learning_rate=learning_rate,
# buffer_size=buffer_size,
# learning_starts=learning_starts,
# batch_size=batch_size,
# tau=tau,
# gamma=gamma,
# train_freq=train_freq,
# gradient_steps=gradient_steps,
# replay_buffer_class=replay_buffer_class, # No action noise
# replay_buffer_kwargs=replay_buffer_kwargs,
# optimize_memory_usage=optimize_memory_usage,
# target_update_interval=target_update_interval,
# exploration_fraction=exploration_fraction,
# exploration_initial_eps=exploration_initial_eps,
# exploration_final_eps=exploration_final_eps,
# max_grad_norm=max_grad_norm,
# tensorboard_log=tensorboard_log,
# create_eval_env=create_eval_env,
# policy_kwargs=policy_kwargs,
# verbose=verbose,
# seed=seed,
# device=device,
# _init_setup_model=_init_setup_model
# )