cleanup models; beef up models to 500M
This commit is contained in:
Dobromir Popov
@ -13,15 +13,13 @@ import torch.nn.functional as F
|
||||
# Add parent directory to path
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||||
|
||||
from NN.models.simple_cnn import CNNModelPyTorch
|
||||
|
||||
# Configure logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class DQNAgent:
|
||||
"""
|
||||
Deep Q-Network agent for trading
|
||||
Uses CNN model as the base network with GPU support
|
||||
Uses Enhanced CNN model as the base network with GPU support for improved performance
|
||||
"""
|
||||
def __init__(self,
|
||||
state_shape: Tuple[int, ...],
|
||||
@ -59,23 +57,18 @@ class DQNAgent:
|
||||
self.batch_size = batch_size
|
||||
self.target_update = target_update
|
||||
|
||||
# Set device for computation (default to CPU)
|
||||
# Set device for computation (default to GPU if available)
|
||||
if device is None:
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
else:
|
||||
self.device = device
|
||||
|
||||
# Initialize models with appropriate architecture based on state shape
|
||||
if isinstance(self.state_dim, tuple) and len(self.state_dim) > 1:
|
||||
# For image-like states (from RL environment with CNN)
|
||||
from NN.models.simple_cnn import SimpleCNN
|
||||
self.policy_net = SimpleCNN(self.state_dim, self.n_actions)
|
||||
self.target_net = SimpleCNN(self.state_dim, self.n_actions)
|
||||
else:
|
||||
# For 1D state vectors (most environments)
|
||||
from NN.models.simple_mlp import SimpleMLP
|
||||
self.policy_net = SimpleMLP(self.state_dim, self.n_actions)
|
||||
self.target_net = SimpleMLP(self.state_dim, self.n_actions)
|
||||
# Initialize models with Enhanced CNN architecture for better performance
|
||||
from NN.models.enhanced_cnn import EnhancedCNN
|
||||
|
||||
# Use Enhanced CNN for both policy and target networks
|
||||
self.policy_net = EnhancedCNN(self.state_dim, self.n_actions)
|
||||
self.target_net = EnhancedCNN(self.state_dim, self.n_actions)
|
||||
|
||||
# Initialize the target network with the same weights as the policy network
|
||||
self.target_net.load_state_dict(self.policy_net.state_dict())
|
||||
@ -166,11 +159,15 @@ class DQNAgent:
|
||||
self.state_size = np.prod(state_shape)
|
||||
self.action_size = n_actions
|
||||
self.memory_size = buffer_size
|
||||
self.timeframes = ["1m", "5m", "15m"][:self.state_dim[0]] # Default timeframes
|
||||
self.timeframes = ["1m", "5m", "15m"][:self.state_dim[0] if isinstance(self.state_dim, tuple) else 3] # Default timeframes
|
||||
|
||||
logger.info(f"DQN Agent using device: {self.device}")
|
||||
logger.info(f"DQN Agent using Enhanced CNN with device: {self.device}")
|
||||
logger.info(f"Trade action fee set to {self.trade_action_fee}, minimum confidence: {self.minimum_action_confidence}")
|
||||
|
||||
# Log model parameters
|
||||
total_params = sum(p.numel() for p in self.policy_net.parameters())
|
||||
logger.info(f"Enhanced CNN Policy Network: {total_params:,} parameters")
|
||||
|
||||
def move_models_to_device(self, device=None):
|
||||
"""Move models to the specified device (GPU/CPU)"""
|
||||
if device is not None:
|
||||
@ -300,7 +297,7 @@ class DQNAgent:
|
||||
|
||||
# Get predictions using the policy network
|
||||
self.policy_net.eval() # Set to evaluation mode for inference
|
||||
action_probs, extrema_pred, price_predictions, hidden_features = self.policy_net(state_tensor)
|
||||
action_probs, extrema_pred, price_predictions, hidden_features, advanced_predictions = self.policy_net(state_tensor)
|
||||
self.policy_net.train() # Back to training mode
|
||||
|
||||
# Store hidden features for integration
|
||||
@ -650,12 +647,12 @@ class DQNAgent:
|
||||
dones = torch.FloatTensor(np.array(dones)).to(self.device)
|
||||
|
||||
# Get current Q values
|
||||
current_q_values, current_extrema_pred, current_price_pred, hidden_features = self.policy_net(states)
|
||||
current_q_values, current_extrema_pred, current_price_pred, hidden_features, current_advanced_pred = self.policy_net(states)
|
||||
current_q_values = current_q_values.gather(1, actions.unsqueeze(1)).squeeze(1)
|
||||
|
||||
# Get next Q values with target network
|
||||
with torch.no_grad():
|
||||
next_q_values, next_extrema_pred, next_price_pred, next_hidden_features = self.target_net(next_states)
|
||||
next_q_values, next_extrema_pred, next_price_pred, next_hidden_features, next_advanced_pred = self.target_net(next_states)
|
||||
next_q_values = next_q_values.max(1)[0]
|
||||
|
||||
# Check for dimension mismatch between rewards and next_q_values
|
||||
@ -727,12 +724,12 @@ class DQNAgent:
|
||||
# Forward pass with amp autocasting
|
||||
with torch.cuda.amp.autocast():
|
||||
# Get current Q values and extrema predictions
|
||||
current_q_values, current_extrema_pred, current_price_pred, hidden_features = self.policy_net(states)
|
||||
current_q_values, current_extrema_pred, current_price_pred, hidden_features, current_advanced_pred = self.policy_net(states)
|
||||
current_q_values = current_q_values.gather(1, actions.unsqueeze(1)).squeeze(1)
|
||||
|
||||
# Get next Q values from target network
|
||||
with torch.no_grad():
|
||||
next_q_values, next_extrema_pred, next_price_pred, next_hidden_features = self.target_net(next_states)
|
||||
next_q_values, next_extrema_pred, next_price_pred, next_hidden_features, next_advanced_pred = self.target_net(next_states)
|
||||
next_q_values = next_q_values.max(1)[0]
|
||||
|
||||
# Check for dimension mismatch and fix it
|
||||
|
||||
@ -110,108 +110,213 @@ class EnhancedCNN(nn.Module):
|
||||
logger.info(f"EnhancedCNN initialized with input shape: {input_shape}, actions: {n_actions}")
|
||||
|
||||
def _build_network(self):
|
||||
"""Build the enhanced neural network with current feature dimensions"""
|
||||
"""Build the MASSIVELY enhanced neural network for 4GB VRAM budget"""
|
||||
|
||||
# 1D CNN for sequential data
|
||||
# MASSIVELY SCALED ARCHITECTURE for 4GB VRAM (up to ~50M parameters)
|
||||
if self.channels > 1:
|
||||
# Reshape expected: [batch, timeframes, features]
|
||||
# Massive convolutional backbone with deeper residual blocks
|
||||
self.conv_layers = nn.Sequential(
|
||||
nn.Conv1d(self.channels, 64, kernel_size=3, padding=1),
|
||||
nn.BatchNorm1d(64),
|
||||
# Initial large conv block
|
||||
nn.Conv1d(self.channels, 256, kernel_size=7, padding=3), # Much wider initial layer
|
||||
nn.BatchNorm1d(256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.1),
|
||||
|
||||
# First residual stage - 256 channels
|
||||
ResidualBlock(256, 512),
|
||||
ResidualBlock(512, 512),
|
||||
ResidualBlock(512, 512),
|
||||
nn.MaxPool1d(kernel_size=2, stride=2),
|
||||
nn.Dropout(0.2),
|
||||
|
||||
ResidualBlock(64, 128),
|
||||
# Second residual stage - 512 channels
|
||||
ResidualBlock(512, 1024),
|
||||
ResidualBlock(1024, 1024),
|
||||
ResidualBlock(1024, 1024),
|
||||
nn.MaxPool1d(kernel_size=2, stride=2),
|
||||
nn.Dropout(0.25),
|
||||
|
||||
# Third residual stage - 1024 channels
|
||||
ResidualBlock(1024, 1536),
|
||||
ResidualBlock(1536, 1536),
|
||||
ResidualBlock(1536, 1536),
|
||||
nn.MaxPool1d(kernel_size=2, stride=2),
|
||||
nn.Dropout(0.3),
|
||||
|
||||
ResidualBlock(128, 256),
|
||||
nn.MaxPool1d(kernel_size=2, stride=2),
|
||||
nn.Dropout(0.4),
|
||||
|
||||
ResidualBlock(256, 512),
|
||||
# Fourth residual stage - 1536 channels (MASSIVE)
|
||||
ResidualBlock(1536, 2048),
|
||||
ResidualBlock(2048, 2048),
|
||||
ResidualBlock(2048, 2048),
|
||||
nn.AdaptiveAvgPool1d(1) # Global average pooling
|
||||
)
|
||||
# Feature dimension after conv layers
|
||||
self.conv_features = 512
|
||||
# Massive feature dimension after conv layers
|
||||
self.conv_features = 2048
|
||||
else:
|
||||
# For 1D vectors, skip the convolutional part
|
||||
# For 1D vectors, use massive dense preprocessing
|
||||
self.conv_layers = None
|
||||
self.conv_features = 0
|
||||
|
||||
# Fully connected layers for all cases
|
||||
# We'll use deeper layers with skip connections
|
||||
# MASSIVE fully connected feature extraction layers
|
||||
if self.conv_layers is None:
|
||||
# For 1D inputs without conv preprocessing
|
||||
self.fc1 = nn.Linear(self.feature_dim, 512)
|
||||
self.features_dim = 512
|
||||
# For 1D inputs - massive feature extraction
|
||||
self.fc1 = nn.Linear(self.feature_dim, 2048)
|
||||
self.features_dim = 2048
|
||||
else:
|
||||
# For data processed by conv layers
|
||||
self.fc1 = nn.Linear(self.conv_features, 512)
|
||||
self.features_dim = 512
|
||||
# For data processed by massive conv layers
|
||||
self.fc1 = nn.Linear(self.conv_features, 2048)
|
||||
self.features_dim = 2048
|
||||
|
||||
# Common feature extraction layers
|
||||
# MASSIVE common feature extraction with multiple attention layers
|
||||
self.fc_layers = nn.Sequential(
|
||||
self.fc1,
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.4),
|
||||
nn.Linear(512, 512),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(2048, 2048), # Keep massive width
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.4),
|
||||
nn.Linear(512, 256),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(2048, 1536), # Still very wide
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(1536, 1024), # Large hidden layer
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(1024, 768), # Final feature representation
|
||||
nn.ReLU()
|
||||
)
|
||||
|
||||
# Dueling architecture
|
||||
# Multiple attention mechanisms for different aspects
|
||||
self.price_attention = SelfAttention(768)
|
||||
self.volume_attention = SelfAttention(768)
|
||||
self.trend_attention = SelfAttention(768)
|
||||
self.volatility_attention = SelfAttention(768)
|
||||
|
||||
# Attention fusion layer
|
||||
self.attention_fusion = nn.Sequential(
|
||||
nn.Linear(768 * 4, 1024), # Combine all attention outputs
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(1024, 768)
|
||||
)
|
||||
|
||||
# MASSIVE dueling architecture with deeper networks
|
||||
self.advantage_stream = nn.Sequential(
|
||||
nn.Linear(768, 512),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(512, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, self.n_actions)
|
||||
)
|
||||
|
||||
self.value_stream = nn.Sequential(
|
||||
nn.Linear(768, 512),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(512, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 1)
|
||||
)
|
||||
|
||||
# Extrema detection head with increased capacity
|
||||
# MASSIVE extrema detection head with ensemble predictions
|
||||
self.extrema_head = nn.Sequential(
|
||||
nn.Linear(256, 128),
|
||||
nn.Linear(768, 512),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(512, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 3) # 0=bottom, 1=top, 2=neither
|
||||
)
|
||||
|
||||
# Price prediction heads with increased capacity
|
||||
# MASSIVE multi-timeframe price prediction heads
|
||||
self.price_pred_immediate = nn.Sequential(
|
||||
nn.Linear(256, 64),
|
||||
nn.Linear(768, 256),
|
||||
nn.ReLU(),
|
||||
nn.Linear(64, 3) # Up, Down, Sideways
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 3) # Up, Down, Sideways
|
||||
)
|
||||
|
||||
self.price_pred_midterm = nn.Sequential(
|
||||
nn.Linear(256, 64),
|
||||
nn.Linear(768, 256),
|
||||
nn.ReLU(),
|
||||
nn.Linear(64, 3) # Up, Down, Sideways
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 3) # Up, Down, Sideways
|
||||
)
|
||||
|
||||
self.price_pred_longterm = nn.Sequential(
|
||||
nn.Linear(256, 64),
|
||||
nn.ReLU(),
|
||||
nn.Linear(64, 3) # Up, Down, Sideways
|
||||
)
|
||||
|
||||
# Value prediction with increased capacity
|
||||
self.price_pred_value = nn.Sequential(
|
||||
nn.Linear(256, 128),
|
||||
nn.Linear(768, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(128, 4) # % change for different timeframes
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 3) # Up, Down, Sideways
|
||||
)
|
||||
|
||||
# Additional attention layer for feature refinement
|
||||
self.attention = SelfAttention(256)
|
||||
# MASSIVE value prediction with ensemble approaches
|
||||
self.price_pred_value = nn.Sequential(
|
||||
nn.Linear(768, 512),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(512, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 8) # More granular % change predictions for different timeframes
|
||||
)
|
||||
|
||||
# Additional specialized prediction heads for better accuracy
|
||||
# Volatility prediction head
|
||||
self.volatility_head = nn.Sequential(
|
||||
nn.Linear(768, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 5) # Very low, low, medium, high, very high volatility
|
||||
)
|
||||
|
||||
# Support/Resistance level detection head
|
||||
self.support_resistance_head = nn.Sequential(
|
||||
nn.Linear(768, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 6) # Strong support, weak support, neutral, weak resistance, strong resistance, breakout
|
||||
)
|
||||
|
||||
# Market regime classification head
|
||||
self.market_regime_head = nn.Sequential(
|
||||
nn.Linear(768, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 7) # Bull trend, bear trend, sideways, volatile up, volatile down, accumulation, distribution
|
||||
)
|
||||
|
||||
# Risk assessment head
|
||||
self.risk_head = nn.Sequential(
|
||||
nn.Linear(768, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 4) # Low risk, medium risk, high risk, extreme risk
|
||||
)
|
||||
|
||||
def _check_rebuild_network(self, features):
|
||||
"""Check if network needs to be rebuilt for different feature dimensions"""
|
||||
@ -225,7 +330,7 @@ class EnhancedCNN(nn.Module):
|
||||
return False
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass through the network"""
|
||||
"""Forward pass through the MASSIVE network"""
|
||||
batch_size = x.size(0)
|
||||
|
||||
# Process different input shapes
|
||||
@ -243,7 +348,7 @@ class EnhancedCNN(nn.Module):
|
||||
total_features = x_reshaped.size(1) * x_reshaped.size(2)
|
||||
self._check_rebuild_network(total_features)
|
||||
|
||||
# Apply convolutions
|
||||
# Apply massive convolutions
|
||||
x_conv = self.conv_layers(x_reshaped)
|
||||
# Flatten: [batch, channels, 1] -> [batch, channels]
|
||||
x_flat = x_conv.view(batch_size, -1)
|
||||
@ -258,31 +363,59 @@ class EnhancedCNN(nn.Module):
|
||||
if x_flat.size(1) != self.feature_dim:
|
||||
self._check_rebuild_network(x_flat.size(1))
|
||||
|
||||
# Apply FC layers
|
||||
features = self.fc_layers(x_flat)
|
||||
# Apply MASSIVE FC layers to get base features
|
||||
features = self.fc_layers(x_flat) # [batch, 768]
|
||||
|
||||
# Add attention for feature refinement
|
||||
features_3d = features.unsqueeze(1) # [batch, 1, features]
|
||||
features_attended, _ = self.attention(features_3d)
|
||||
features_refined = features_attended.squeeze(1) # [batch, features]
|
||||
# Apply multiple specialized attention mechanisms
|
||||
features_3d = features.unsqueeze(1) # [batch, 1, 768]
|
||||
|
||||
# Calculate advantage and value
|
||||
# Get attention-refined features for different aspects
|
||||
price_features, _ = self.price_attention(features_3d)
|
||||
price_features = price_features.squeeze(1) # [batch, 768]
|
||||
|
||||
volume_features, _ = self.volume_attention(features_3d)
|
||||
volume_features = volume_features.squeeze(1) # [batch, 768]
|
||||
|
||||
trend_features, _ = self.trend_attention(features_3d)
|
||||
trend_features = trend_features.squeeze(1) # [batch, 768]
|
||||
|
||||
volatility_features, _ = self.volatility_attention(features_3d)
|
||||
volatility_features = volatility_features.squeeze(1) # [batch, 768]
|
||||
|
||||
# Fuse all attention outputs
|
||||
combined_attention = torch.cat([
|
||||
price_features, volume_features,
|
||||
trend_features, volatility_features
|
||||
], dim=1) # [batch, 768*4]
|
||||
|
||||
# Apply attention fusion to get final refined features
|
||||
features_refined = self.attention_fusion(combined_attention) # [batch, 768]
|
||||
|
||||
# Calculate advantage and value (Dueling DQN architecture)
|
||||
advantage = self.advantage_stream(features_refined)
|
||||
value = self.value_stream(features_refined)
|
||||
|
||||
# Combine for Q-values (Dueling architecture)
|
||||
q_values = value + advantage - advantage.mean(dim=1, keepdim=True)
|
||||
|
||||
# Get extrema predictions
|
||||
# Get massive ensemble of predictions
|
||||
|
||||
# Extrema predictions (bottom/top/neither detection)
|
||||
extrema_pred = self.extrema_head(features_refined)
|
||||
|
||||
# Price movement predictions
|
||||
# Multi-timeframe price movement predictions
|
||||
price_immediate = self.price_pred_immediate(features_refined)
|
||||
price_midterm = self.price_pred_midterm(features_refined)
|
||||
price_longterm = self.price_pred_longterm(features_refined)
|
||||
price_values = self.price_pred_value(features_refined)
|
||||
|
||||
# Package price predictions
|
||||
# Additional specialized predictions for enhanced accuracy
|
||||
volatility_pred = self.volatility_head(features_refined)
|
||||
support_resistance_pred = self.support_resistance_head(features_refined)
|
||||
market_regime_pred = self.market_regime_head(features_refined)
|
||||
risk_pred = self.risk_head(features_refined)
|
||||
|
||||
# Package all price predictions
|
||||
price_predictions = {
|
||||
'immediate': price_immediate,
|
||||
'midterm': price_midterm,
|
||||
@ -290,31 +423,60 @@ class EnhancedCNN(nn.Module):
|
||||
'values': price_values
|
||||
}
|
||||
|
||||
return q_values, extrema_pred, price_predictions, features_refined
|
||||
# Package additional predictions for enhanced decision making
|
||||
advanced_predictions = {
|
||||
'volatility': volatility_pred,
|
||||
'support_resistance': support_resistance_pred,
|
||||
'market_regime': market_regime_pred,
|
||||
'risk_assessment': risk_pred
|
||||
}
|
||||
|
||||
return q_values, extrema_pred, price_predictions, features_refined, advanced_predictions
|
||||
|
||||
def act(self, state, explore=True):
|
||||
"""
|
||||
Choose action based on state with confidence thresholding
|
||||
"""
|
||||
"""Enhanced action selection with massive model predictions"""
|
||||
if explore and np.random.random() < 0.1: # 10% random exploration
|
||||
return np.random.choice(self.n_actions)
|
||||
|
||||
self.eval()
|
||||
state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
q_values, _, _, _ = self(state_tensor)
|
||||
q_values, extrema_pred, price_predictions, features, advanced_predictions = self(state_tensor)
|
||||
|
||||
# Apply softmax to get action probabilities
|
||||
action_probs = F.softmax(q_values, dim=1)
|
||||
action_probs = torch.softmax(q_values, dim=1)
|
||||
action = torch.argmax(action_probs, dim=1).item()
|
||||
|
||||
# Get action with highest probability
|
||||
action = action_probs.argmax(dim=1).item()
|
||||
action_confidence = action_probs[0, action].item()
|
||||
# Log advanced predictions for better decision making
|
||||
if hasattr(self, '_log_predictions') and self._log_predictions:
|
||||
# Log volatility prediction
|
||||
volatility = torch.softmax(advanced_predictions['volatility'], dim=1)
|
||||
volatility_class = torch.argmax(volatility, dim=1).item()
|
||||
volatility_labels = ['Very Low', 'Low', 'Medium', 'High', 'Very High']
|
||||
|
||||
# Log support/resistance prediction
|
||||
sr = torch.softmax(advanced_predictions['support_resistance'], dim=1)
|
||||
sr_class = torch.argmax(sr, dim=1).item()
|
||||
sr_labels = ['Strong Support', 'Weak Support', 'Neutral', 'Weak Resistance', 'Strong Resistance', 'Breakout']
|
||||
|
||||
# Log market regime prediction
|
||||
regime = torch.softmax(advanced_predictions['market_regime'], dim=1)
|
||||
regime_class = torch.argmax(regime, dim=1).item()
|
||||
regime_labels = ['Bull Trend', 'Bear Trend', 'Sideways', 'Volatile Up', 'Volatile Down', 'Accumulation', 'Distribution']
|
||||
|
||||
# Log risk assessment
|
||||
risk = torch.softmax(advanced_predictions['risk_assessment'], dim=1)
|
||||
risk_class = torch.argmax(risk, dim=1).item()
|
||||
risk_labels = ['Low Risk', 'Medium Risk', 'High Risk', 'Extreme Risk']
|
||||
|
||||
logger.info(f"MASSIVE Model Predictions:")
|
||||
logger.info(f" Volatility: {volatility_labels[volatility_class]} ({volatility[0, volatility_class]:.3f})")
|
||||
logger.info(f" Support/Resistance: {sr_labels[sr_class]} ({sr[0, sr_class]:.3f})")
|
||||
logger.info(f" Market Regime: {regime_labels[regime_class]} ({regime[0, regime_class]:.3f})")
|
||||
logger.info(f" Risk Level: {risk_labels[risk_class]} ({risk[0, risk_class]:.3f})")
|
||||
|
||||
# Check if confidence exceeds threshold
|
||||
if action_confidence < self.confidence_threshold:
|
||||
# Force HOLD action (typically action 2)
|
||||
action = 2 # Assume 2 is HOLD
|
||||
logger.info(f"Action {action} confidence {action_confidence:.4f} below threshold {self.confidence_threshold}, forcing HOLD")
|
||||
|
||||
return action, action_confidence
|
||||
return action
|
||||
|
||||
def save(self, path):
|
||||
"""Save model weights and architecture"""
|
||||
|
||||
@ -1,500 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
import numpy as np
|
||||
import os
|
||||
import logging
|
||||
import torch.nn.functional as F
|
||||
from typing import List, Tuple
|
||||
|
||||
# Configure logger
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class PricePatternAttention(nn.Module):
|
||||
"""
|
||||
Attention mechanism specifically designed to focus on price patterns
|
||||
that might indicate local extrema or trend reversals
|
||||
"""
|
||||
def __init__(self, input_dim, hidden_dim=64):
|
||||
super(PricePatternAttention, self).__init__()
|
||||
self.query = nn.Linear(input_dim, hidden_dim)
|
||||
self.key = nn.Linear(input_dim, hidden_dim)
|
||||
self.value = nn.Linear(input_dim, hidden_dim)
|
||||
self.scale = torch.sqrt(torch.tensor(hidden_dim, dtype=torch.float32))
|
||||
|
||||
def forward(self, x):
|
||||
"""Apply attention to input sequence"""
|
||||
# x shape: [batch_size, seq_len, features]
|
||||
batch_size, seq_len, _ = x.size()
|
||||
|
||||
# Project input to query, key, value
|
||||
q = self.query(x) # [batch_size, seq_len, hidden_dim]
|
||||
k = self.key(x) # [batch_size, seq_len, hidden_dim]
|
||||
v = self.value(x) # [batch_size, seq_len, hidden_dim]
|
||||
|
||||
# Calculate attention scores
|
||||
scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale # [batch_size, seq_len, seq_len]
|
||||
|
||||
# Apply softmax to get attention weights
|
||||
attn_weights = F.softmax(scores, dim=-1) # [batch_size, seq_len, seq_len]
|
||||
|
||||
# Apply attention to values
|
||||
output = torch.matmul(attn_weights, v) # [batch_size, seq_len, hidden_dim]
|
||||
|
||||
return output, attn_weights
|
||||
|
||||
class AdaptiveNorm(nn.Module):
|
||||
"""
|
||||
Adaptive normalization layer that chooses between different normalization
|
||||
methods based on input dimensions
|
||||
"""
|
||||
def __init__(self, num_features):
|
||||
super(AdaptiveNorm, self).__init__()
|
||||
self.batch_norm = nn.BatchNorm1d(num_features, affine=True)
|
||||
self.group_norm = nn.GroupNorm(min(32, num_features), num_features)
|
||||
self.layer_norm = nn.LayerNorm([num_features, 1])
|
||||
|
||||
def forward(self, x):
|
||||
# Check input dimensions
|
||||
batch_size, channels, seq_len = x.size()
|
||||
|
||||
# Choose normalization method:
|
||||
# - Batch size > 1 and seq_len > 1: BatchNorm
|
||||
# - Batch size == 1 or seq_len == 1: GroupNorm
|
||||
# - Fallback for extreme cases: LayerNorm
|
||||
if batch_size > 1 and seq_len > 1:
|
||||
return self.batch_norm(x)
|
||||
elif seq_len > 1:
|
||||
return self.group_norm(x)
|
||||
else:
|
||||
# For 1D inputs (seq_len=1), we need to adjust the layer norm
|
||||
# to the actual input size
|
||||
if not hasattr(self, 'layer_norm_1d') or self.layer_norm_1d.normalized_shape[0] != channels:
|
||||
self.layer_norm_1d = nn.LayerNorm([channels, seq_len]).to(x.device)
|
||||
return self.layer_norm_1d(x)
|
||||
|
||||
class SimpleCNN(nn.Module):
|
||||
"""
|
||||
Simple CNN model for reinforcement learning with image-like state inputs
|
||||
"""
|
||||
def __init__(self, input_shape, n_actions):
|
||||
super(SimpleCNN, self).__init__()
|
||||
|
||||
# Store dimensions
|
||||
self.input_shape = input_shape
|
||||
self.n_actions = n_actions
|
||||
|
||||
# Calculate input dimensions
|
||||
if len(input_shape) == 3: # [channels, height, width]
|
||||
self.channels, self.height, self.width = input_shape
|
||||
self.feature_dim = self.height * self.width
|
||||
elif len(input_shape) == 2: # [timeframes, features]
|
||||
self.channels = input_shape[0]
|
||||
self.features = input_shape[1]
|
||||
self.feature_dim = self.features
|
||||
elif len(input_shape) == 1: # [features]
|
||||
self.channels = 1
|
||||
self.features = input_shape[0]
|
||||
self.feature_dim = self.features
|
||||
else:
|
||||
raise ValueError(f"Unsupported input shape: {input_shape}")
|
||||
|
||||
# Build network
|
||||
self._build_network()
|
||||
|
||||
# Initialize device
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
self.to(self.device)
|
||||
|
||||
logger.info(f"SimpleCNN initialized with input shape: {input_shape}, actions: {n_actions}")
|
||||
|
||||
def _build_network(self):
|
||||
"""Build the neural network with current feature dimensions"""
|
||||
# Create a flexible architecture that adapts to input dimensions
|
||||
# Increased complexity
|
||||
self.fc_layers = nn.Sequential(
|
||||
nn.Linear(self.feature_dim, 512), # Increased size
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.2), # Added dropout
|
||||
nn.Linear(512, 512), # Increased size
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.2), # Added dropout
|
||||
nn.Linear(512, 512), # Added layer
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.2) # Added dropout
|
||||
)
|
||||
|
||||
# Output heads (Dueling DQN architecture)
|
||||
self.advantage_head = nn.Linear(512, self.n_actions) # Updated input size
|
||||
self.value_head = nn.Linear(512, 1) # Updated input size
|
||||
|
||||
# Extrema detection head
|
||||
self.extrema_head = nn.Linear(512, 3) # 0=bottom, 1=top, 2=neither, Updated input size
|
||||
|
||||
# Price prediction heads for different timeframes
|
||||
self.price_pred_immediate = nn.Linear(512, 3) # Updated input size
|
||||
self.price_pred_midterm = nn.Linear(512, 3) # Updated input size
|
||||
self.price_pred_longterm = nn.Linear(512, 3) # Updated input size
|
||||
|
||||
# Regression heads for exact price prediction
|
||||
self.price_pred_value = nn.Linear(512, 4) # Updated input size
|
||||
|
||||
def _check_rebuild_network(self, features):
|
||||
"""Check if network needs to be rebuilt for different feature dimensions"""
|
||||
if features != self.feature_dim:
|
||||
logger.info(f"Rebuilding network for new feature dimension: {features} (was {self.feature_dim})")
|
||||
self.feature_dim = features
|
||||
self._build_network()
|
||||
# Move to device after rebuilding
|
||||
self.to(self.device)
|
||||
return True
|
||||
return False
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass through the network"""
|
||||
# Flatten input if needed to ensure it matches the expected feature dimension
|
||||
batch_size = x.size(0)
|
||||
|
||||
# Reshape input if needed
|
||||
if len(x.shape) > 2: # Handle multi-dimensional input
|
||||
# For 3D input: [batch, seq_len, features] or [batch, channels, features]
|
||||
x = x.reshape(batch_size, -1) # Flatten to [batch, seq_len*features]
|
||||
|
||||
# Check if the feature dimension matches and rebuild if necessary
|
||||
if x.size(1) != self.feature_dim:
|
||||
self._check_rebuild_network(x.size(1))
|
||||
|
||||
# Apply fully connected layers with ReLU activation
|
||||
x = self.fc_layers(x)
|
||||
|
||||
# Branch 1: Action values (Q-values)
|
||||
action_values = self.advantage_head(x)
|
||||
|
||||
# Branch 2: Extrema detection (market top/bottom classification)
|
||||
extrema_pred = self.extrema_head(x)
|
||||
|
||||
# Branch 3: Price movement prediction over different timeframes
|
||||
# Split into three timeframes: immediate, midterm, longterm
|
||||
price_immediate = self.price_pred_immediate(x)
|
||||
price_midterm = self.price_pred_midterm(x)
|
||||
price_longterm = self.price_pred_longterm(x)
|
||||
|
||||
# Branch 4: Value prediction (regression for expected price changes)
|
||||
price_values = self.price_pred_value(x)
|
||||
|
||||
# Package price predictions
|
||||
price_predictions = {
|
||||
'immediate': price_immediate, # Classification (up/down/sideways)
|
||||
'midterm': price_midterm, # Classification (up/down/sideways)
|
||||
'longterm': price_longterm, # Classification (up/down/sideways)
|
||||
'values': price_values # Regression (expected % change)
|
||||
}
|
||||
|
||||
# Return all outputs and the hidden feature representation
|
||||
return action_values, extrema_pred, price_predictions, x
|
||||
|
||||
def extract_features(self, x):
|
||||
"""Extract hidden features from the input and return both action values and features"""
|
||||
# Flatten input if needed to ensure it matches the expected feature dimension
|
||||
batch_size = x.size(0)
|
||||
|
||||
# Reshape input if needed
|
||||
if len(x.shape) > 2: # Handle multi-dimensional input
|
||||
# For 3D input: [batch, seq_len, features] or [batch, channels, features]
|
||||
x = x.reshape(batch_size, -1) # Flatten to [batch, seq_len*features]
|
||||
|
||||
# Check if the feature dimension matches and rebuild if necessary
|
||||
if x.size(1) != self.feature_dim:
|
||||
self._check_rebuild_network(x.size(1))
|
||||
|
||||
# Apply fully connected layers with ReLU activation
|
||||
x_features = self.fc_layers(x)
|
||||
|
||||
# Branch 1: Action values (Q-values)
|
||||
action_values = self.advantage_head(x_features)
|
||||
|
||||
# Return action values and the hidden feature representation
|
||||
return action_values, x_features
|
||||
|
||||
def save(self, path):
|
||||
"""Save model weights and architecture"""
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
torch.save({
|
||||
'state_dict': self.state_dict(),
|
||||
'input_shape': self.input_shape,
|
||||
'n_actions': self.n_actions,
|
||||
'feature_dim': self.feature_dim
|
||||
}, f"{path}.pt")
|
||||
logger.info(f"Model saved to {path}.pt")
|
||||
|
||||
def load(self, path):
|
||||
"""Load model weights and architecture"""
|
||||
try:
|
||||
checkpoint = torch.load(f"{path}.pt", map_location=self.device)
|
||||
self.input_shape = checkpoint['input_shape']
|
||||
self.n_actions = checkpoint['n_actions']
|
||||
self.feature_dim = checkpoint['feature_dim']
|
||||
self._build_network()
|
||||
self.load_state_dict(checkpoint['state_dict'])
|
||||
self.to(self.device)
|
||||
logger.info(f"Model loaded from {path}.pt")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading model: {str(e)}")
|
||||
return False
|
||||
|
||||
class CNNModelPyTorch(nn.Module):
|
||||
"""
|
||||
CNN model for trading with multiple timeframes
|
||||
"""
|
||||
def __init__(self, window_size=20, num_features=5, output_size=3, timeframes=None):
|
||||
super(CNNModelPyTorch, self).__init__()
|
||||
|
||||
if timeframes is None:
|
||||
timeframes = [1]
|
||||
|
||||
self.window_size = window_size
|
||||
self.num_features = num_features
|
||||
self.output_size = output_size
|
||||
self.timeframes = timeframes
|
||||
|
||||
# num_features should already be the total features across all timeframes
|
||||
self.total_features = num_features
|
||||
logger.info(f"CNNModelPyTorch initialized with window_size={window_size}, num_features={num_features}, "
|
||||
f"total_features={self.total_features}, output_size={output_size}, timeframes={timeframes}")
|
||||
|
||||
# Device configuration
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
logger.info(f"Using device: {self.device}")
|
||||
|
||||
# Create model architecture
|
||||
self._create_layers()
|
||||
|
||||
# Move model to device
|
||||
self.to(self.device)
|
||||
|
||||
def _create_layers(self):
|
||||
"""Create all model layers with current feature dimensions"""
|
||||
# Convolutional layers - use total_features as input channels
|
||||
self.conv1 = nn.Conv1d(self.total_features, 64, kernel_size=3, padding=1)
|
||||
self.norm1 = AdaptiveNorm(64)
|
||||
self.dropout1 = nn.Dropout(0.2)
|
||||
|
||||
self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
|
||||
self.norm2 = AdaptiveNorm(128)
|
||||
self.dropout2 = nn.Dropout(0.3)
|
||||
|
||||
self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
|
||||
self.norm3 = AdaptiveNorm(256)
|
||||
self.dropout3 = nn.Dropout(0.4)
|
||||
|
||||
# Add price pattern attention layer
|
||||
self.attention = PricePatternAttention(256)
|
||||
|
||||
# Extrema detection specialized convolutional layer
|
||||
self.extrema_conv = nn.Conv1d(256, 128, kernel_size=3, padding=1) # Smaller kernel for small inputs
|
||||
self.extrema_norm = AdaptiveNorm(128)
|
||||
|
||||
# Fully connected layers - input size will be determined dynamically
|
||||
self.fc1 = None # Will be initialized in forward pass
|
||||
self.fc2 = nn.Linear(512, 256)
|
||||
self.dropout_fc = nn.Dropout(0.5)
|
||||
|
||||
# Advantage and Value streams (Dueling DQN architecture)
|
||||
self.fc3 = nn.Linear(256, self.output_size) # Advantage stream
|
||||
self.value_fc = nn.Linear(256, 1) # Value stream
|
||||
|
||||
# Additional prediction head for extrema detection (tops/bottoms)
|
||||
self.extrema_fc = nn.Linear(256, 3) # 0=bottom, 1=top, 2=neither
|
||||
|
||||
# Initialize optimizer and scheduler
|
||||
self.optimizer = optim.Adam(self.parameters(), lr=0.001)
|
||||
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
|
||||
self.optimizer, mode='max', factor=0.5, patience=5, verbose=True
|
||||
)
|
||||
|
||||
def rebuild_conv_layers(self, input_channels):
|
||||
"""
|
||||
Rebuild convolutional layers for different input dimensions
|
||||
|
||||
Args:
|
||||
input_channels: Number of input channels (features) in the data
|
||||
"""
|
||||
logger.info(f"Rebuilding convolutional layers for {input_channels} input channels")
|
||||
|
||||
# Update total features
|
||||
self.total_features = input_channels
|
||||
|
||||
# Recreate all layers with new dimensions
|
||||
self._create_layers()
|
||||
|
||||
# Move layers to device
|
||||
self.to(self.device)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Forward pass through the network"""
|
||||
# Ensure input is on the correct device
|
||||
x = x.to(self.device)
|
||||
|
||||
# Log input tensor shape for debugging
|
||||
input_shape = x.size()
|
||||
logger.debug(f"Input tensor shape: {input_shape}")
|
||||
|
||||
# Check input dimensions and reshape as needed
|
||||
if len(x.size()) == 2:
|
||||
# If input is [batch_size, features], reshape to [batch_size, features, 1]
|
||||
batch_size, feature_dim = x.size()
|
||||
|
||||
# Check and handle if input features don't match model expectations
|
||||
if feature_dim != self.total_features:
|
||||
logger.warning(f"Input features ({feature_dim}) don't match model features ({self.total_features})")
|
||||
if not hasattr(self, 'rebuild_warning_shown'):
|
||||
logger.error(f"Dimension mismatch: Expected {self.total_features} features but got {feature_dim}")
|
||||
self.rebuild_warning_shown = True
|
||||
# Don't rebuild - instead adapt the input
|
||||
# If features are fewer, pad with zeros. If more, truncate
|
||||
if feature_dim < self.total_features:
|
||||
padding = torch.zeros(batch_size, self.total_features - feature_dim, device=self.device)
|
||||
x = torch.cat([x, padding], dim=1)
|
||||
else:
|
||||
x = x[:, :self.total_features]
|
||||
|
||||
# For 1D input, use a sequence length of 1
|
||||
seq_len = 1
|
||||
x = x.unsqueeze(2) # Reshape to [batch, features, 1]
|
||||
elif len(x.size()) == 3:
|
||||
# Standard case: [batch_size, window_size, features]
|
||||
batch_size, seq_len, feature_dim = x.size()
|
||||
|
||||
# Check and handle if input dimensions don't match model expectations
|
||||
if feature_dim != self.total_features:
|
||||
logger.warning(f"Input features ({feature_dim}) don't match model features ({self.total_features})")
|
||||
if not hasattr(self, 'rebuild_warning_shown'):
|
||||
logger.error(f"Dimension mismatch: Expected {self.total_features} features but got {feature_dim}")
|
||||
self.rebuild_warning_shown = True
|
||||
# Don't rebuild - instead adapt the input
|
||||
# If features are fewer, pad with zeros. If more, truncate
|
||||
if feature_dim < self.total_features:
|
||||
padding = torch.zeros(batch_size, seq_len, self.total_features - feature_dim, device=self.device)
|
||||
x = torch.cat([x, padding], dim=2)
|
||||
else:
|
||||
x = x[:, :, :self.total_features]
|
||||
|
||||
# Reshape input: [batch, window_size, features] -> [batch, features, window_size]
|
||||
x = x.permute(0, 2, 1)
|
||||
else:
|
||||
raise ValueError(f"Unexpected input shape: {x.size()}, expected 2D or 3D tensor")
|
||||
|
||||
# Log reshaped tensor for debugging
|
||||
logger.debug(f"Reshaped tensor for convolution: {x.size()}")
|
||||
|
||||
# Convolutional layers with dropout - safely handle small spatial dimensions
|
||||
try:
|
||||
x = self.dropout1(F.relu(self.norm1(self.conv1(x))))
|
||||
x = self.dropout2(F.relu(self.norm2(self.conv2(x))))
|
||||
x = self.dropout3(F.relu(self.norm3(self.conv3(x))))
|
||||
except Exception as e:
|
||||
logger.warning(f"Error in convolutional layers: {str(e)}")
|
||||
# Fallback for very small inputs: skip some convolutions
|
||||
if seq_len < 3:
|
||||
# Apply a simpler convolution for very small inputs
|
||||
x = F.relu(self.conv1(x))
|
||||
x = F.relu(self.conv2(x))
|
||||
# Skip last conv if we get dimension errors
|
||||
try:
|
||||
x = F.relu(self.conv3(x))
|
||||
except:
|
||||
pass
|
||||
|
||||
# Store conv features for extrema detection
|
||||
conv_features = x
|
||||
|
||||
# Get the current shape after convolutions
|
||||
_, channels, conv_seq_len = x.size()
|
||||
|
||||
# Initialize fc1 if not created yet or if the shape has changed
|
||||
if self.fc1 is None:
|
||||
flattened_size = channels * conv_seq_len
|
||||
logger.info(f"Initializing fc1 with input size {flattened_size}")
|
||||
self.fc1 = nn.Linear(flattened_size, 512).to(self.device)
|
||||
|
||||
# Apply extrema detection safely
|
||||
try:
|
||||
extrema_features = F.relu(self.extrema_norm(self.extrema_conv(conv_features)))
|
||||
except Exception as e:
|
||||
logger.warning(f"Error in extrema detection: {str(e)}")
|
||||
extrema_features = conv_features # Fallback
|
||||
|
||||
# Handle attention for small sequence lengths
|
||||
if conv_seq_len > 1:
|
||||
# Reshape for attention: [batch, channels, seq_len] -> [batch, seq_len, channels]
|
||||
x_attention = x.permute(0, 2, 1)
|
||||
|
||||
# Apply attention
|
||||
try:
|
||||
attention_output, attention_weights = self.attention(x_attention)
|
||||
except Exception as e:
|
||||
logger.warning(f"Error in attention layer: {str(e)}")
|
||||
# Fallback: don't use attention
|
||||
|
||||
# Flatten - get the actual shape for this batch
|
||||
flattened_size = channels * conv_seq_len
|
||||
x = x.view(batch_size, flattened_size)
|
||||
|
||||
# Check if we need to recreate fc1 with the correct size
|
||||
if self.fc1.in_features != flattened_size:
|
||||
logger.info(f"Recreating fc1 layer to match input size {flattened_size}")
|
||||
self.fc1 = nn.Linear(flattened_size, 512).to(self.device)
|
||||
# Reinitialize optimizer after changing the model
|
||||
self.optimizer = optim.Adam(self.parameters(), lr=0.001)
|
||||
|
||||
# Fully connected layers with dropout
|
||||
x = F.relu(self.fc1(x))
|
||||
x = self.dropout_fc(F.relu(self.fc2(x)))
|
||||
|
||||
# Split into advantage and value streams
|
||||
advantage = self.fc3(x)
|
||||
value = self.value_fc(x)
|
||||
|
||||
# Combine value and advantage
|
||||
q_values = value + (advantage - advantage.mean(dim=1, keepdim=True))
|
||||
|
||||
# Also compute extrema prediction from the same features
|
||||
extrema_flat = extrema_features.view(batch_size, -1)
|
||||
extrema_pred = self.extrema_fc(x) # Use the same features for extrema prediction
|
||||
|
||||
return q_values, extrema_pred
|
||||
|
||||
def predict(self, X):
|
||||
"""Make predictions"""
|
||||
self.eval()
|
||||
|
||||
# Convert to tensor if not already
|
||||
if not isinstance(X, torch.Tensor):
|
||||
X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device)
|
||||
else:
|
||||
X_tensor = X.to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
q_values, extrema_pred = self(X_tensor)
|
||||
q_values_np = q_values.cpu().numpy()
|
||||
actions = np.argmax(q_values_np, axis=1)
|
||||
|
||||
# Also return extrema predictions
|
||||
extrema_np = extrema_pred.cpu().numpy()
|
||||
extrema_classes = np.argmax(extrema_np, axis=1)
|
||||
|
||||
return actions, q_values_np, extrema_classes
|
||||
|
||||
def save(self, path: str):
|
||||
"""Save model weights"""
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
torch.save(self.state_dict(), f"{path}.pt")
|
||||
logger.info(f"Model saved to {path}.pt")
|
||||
|
||||
def load(self, path: str):
|
||||
"""Load model weights"""
|
||||
self.load_state_dict(torch.load(f"{path}.pt", map_location=self.device))
|
||||
self.eval()
|
||||
logger.info(f"Model loaded from {path}.pt")
|
||||
@ -1,70 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import os
|
||||
import logging
|
||||
|
||||
# Configure logger
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class SimpleMLP(nn.Module):
|
||||
"""
|
||||
Simple Multi-Layer Perceptron for reinforcement learning with vector state inputs
|
||||
Implements dueling architecture for better Q-learning
|
||||
"""
|
||||
def __init__(self, state_dim, n_actions):
|
||||
super(SimpleMLP, self).__init__()
|
||||
|
||||
# Store dimensions
|
||||
self.state_dim = state_dim
|
||||
self.n_actions = n_actions
|
||||
|
||||
# Calculate input size
|
||||
if isinstance(state_dim, tuple):
|
||||
self.input_size = int(np.prod(state_dim))
|
||||
else:
|
||||
self.input_size = state_dim
|
||||
|
||||
# Hidden layers
|
||||
self.fc1 = nn.Linear(self.input_size, 256)
|
||||
self.fc2 = nn.Linear(256, 256)
|
||||
|
||||
# Dueling architecture
|
||||
self.advantage = nn.Linear(256, n_actions)
|
||||
self.value = nn.Linear(256, 1)
|
||||
|
||||
# Extrema detection
|
||||
self.extrema_head = nn.Linear(256, 3) # 0=bottom, 1=top, 2=neither
|
||||
|
||||
# Move to appropriate device
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
self.to(self.device)
|
||||
|
||||
logger.info(f"SimpleMLP initialized with input size: {self.input_size}, actions: {n_actions}")
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass through the network
|
||||
Returns both action values and extrema predictions
|
||||
"""
|
||||
# Handle different input shapes
|
||||
if isinstance(self.state_dim, tuple) and len(self.state_dim) > 1:
|
||||
x = x.view(-1, self.input_size)
|
||||
|
||||
# Main network
|
||||
x = F.relu(self.fc1(x))
|
||||
x = F.relu(self.fc2(x))
|
||||
|
||||
# Dueling architecture
|
||||
advantage = self.advantage(x)
|
||||
value = self.value(x)
|
||||
|
||||
# Combine value and advantage (Q = V + A - mean(A))
|
||||
q_values = value + advantage - advantage.mean(dim=1, keepdim=True)
|
||||
|
||||
# Extrema predictions
|
||||
extrema = F.softmax(self.extrema_head(x), dim=1)
|
||||
|
||||
return q_values, extrema
|
||||
Reference in New Issue
Block a user