50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
# training/rl_agent.py
|
|
import random
|
|
from collections import deque
|
|
import numpy as np
|
|
import torch
|
|
|
|
class ContinuousRLAgent:
|
|
def __init__(self, model, optimizer, replay_buffer, batch_size=32, gamma=0.99):
|
|
self.model = model
|
|
self.optimizer = optimizer
|
|
self.replay_buffer = replay_buffer
|
|
self.batch_size = batch_size
|
|
self.gamma = gamma
|
|
|
|
def act(self, state, epsilon=0.0):
|
|
"""
|
|
Select an action based on the state, using an epsilon-greedy policy.
|
|
"""
|
|
if random.random() < epsilon:
|
|
# Exploration: choose a random action.
|
|
action = np.random.choice([0, 1, 2]) # SELL, HOLD, BUY
|
|
else:
|
|
# Exploitation: choose the action with the highest Q-value.
|
|
with torch.no_grad():
|
|
state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
|
|
q_values = self.model(state_tensor)
|
|
action = torch.argmax(q_values).item()
|
|
return action
|
|
|
|
class ReplayBuffer:
|
|
def __init__(self, capacity):
|
|
self.buffer = deque(maxlen=capacity)
|
|
|
|
def push(self, state, action, reward, next_state, done):
|
|
"""
|
|
Store an experience tuple into the replay buffer.
|
|
"""
|
|
self.buffer.append((state, action, reward, next_state, done))
|
|
|
|
def sample(self, batch_size):
|
|
"""
|
|
Randomly sample a batch of experiences from the replay buffer.
|
|
"""
|
|
batch = random.sample(self.buffer, batch_size)
|
|
states, actions, rewards, next_states, dones = zip(*batch)
|
|
return states, actions, rewards, next_states, dones
|
|
|
|
def __len__(self):
|
|
return len(self.buffer)
|