272 lines
8.4 KiB
Rust
272 lines
8.4 KiB
Rust
use crate::Result;
|
|
use rtx_tensor::{Device, Tensor};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SACConfig {
|
|
pub learning_rate: f64,
|
|
pub gamma: f32,
|
|
pub tau: f32,
|
|
pub alpha: f32,
|
|
pub target_update_interval: usize,
|
|
pub automatic_entropy_tuning: bool,
|
|
pub target_entropy: Option<f32>,
|
|
pub replay_buffer_size: usize,
|
|
pub batch_size: usize,
|
|
}
|
|
|
|
impl Default for SACConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
learning_rate: 3e-4,
|
|
gamma: 0.99,
|
|
tau: 0.005,
|
|
alpha: 0.2,
|
|
target_update_interval: 1,
|
|
automatic_entropy_tuning: true,
|
|
target_entropy: None,
|
|
replay_buffer_size: 1_000_000,
|
|
batch_size: 256,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct SACMetrics {
|
|
pub actor_loss: f32,
|
|
pub critic1_loss: f32,
|
|
pub critic2_loss: f32,
|
|
pub temperature_loss: f32,
|
|
pub alpha: f32,
|
|
pub target_entropy: f32,
|
|
}
|
|
|
|
pub struct SAC {
|
|
config: SACConfig,
|
|
state_dim: usize,
|
|
action_dim: usize,
|
|
hidden_dim: usize,
|
|
device: Device,
|
|
// Neural network parameters - dummy implementation for testing
|
|
actor_params: Vec<Tensor>,
|
|
critic1_params: Vec<Tensor>,
|
|
critic2_params: Vec<Tensor>,
|
|
target_critic1_params: Vec<Tensor>,
|
|
target_critic2_params: Vec<Tensor>,
|
|
log_alpha: Tensor,
|
|
}
|
|
|
|
impl SAC {
|
|
pub fn new(
|
|
config: SACConfig,
|
|
state_dim: usize,
|
|
action_dim: usize,
|
|
hidden_dim: usize,
|
|
device: Device,
|
|
) -> Self {
|
|
// Initialize dummy parameters for testing
|
|
let actor_params = vec![
|
|
Tensor::randn(&[hidden_dim, state_dim], &device).unwrap(),
|
|
Tensor::randn(&[hidden_dim], &device).unwrap(),
|
|
Tensor::randn(&[action_dim * 2, hidden_dim], &device).unwrap(), // mean and log_std
|
|
Tensor::randn(&[action_dim * 2], &device).unwrap(),
|
|
];
|
|
|
|
let critic1_params = vec![
|
|
Tensor::randn(&[hidden_dim, state_dim + action_dim], &device).unwrap(),
|
|
Tensor::randn(&[hidden_dim], &device).unwrap(),
|
|
Tensor::randn(&[1, hidden_dim], &device).unwrap(),
|
|
Tensor::randn(&[1], &device).unwrap(),
|
|
];
|
|
|
|
let critic2_params = vec![
|
|
Tensor::randn(&[hidden_dim, state_dim + action_dim], &device).unwrap(),
|
|
Tensor::randn(&[hidden_dim], &device).unwrap(),
|
|
Tensor::randn(&[1, hidden_dim], &device).unwrap(),
|
|
Tensor::randn(&[1], &device).unwrap(),
|
|
];
|
|
|
|
let target_critic1_params = critic1_params.clone();
|
|
let target_critic2_params = critic2_params.clone();
|
|
|
|
let log_alpha = Tensor::zeros([1], &device).unwrap();
|
|
|
|
Self {
|
|
config,
|
|
state_dim,
|
|
action_dim,
|
|
hidden_dim,
|
|
device,
|
|
actor_params,
|
|
critic1_params,
|
|
critic2_params,
|
|
target_critic1_params,
|
|
target_critic2_params,
|
|
log_alpha,
|
|
}
|
|
}
|
|
|
|
pub fn state_dim(&self) -> usize {
|
|
self.state_dim
|
|
}
|
|
|
|
pub fn action_dim(&self) -> usize {
|
|
self.action_dim
|
|
}
|
|
|
|
pub async fn actor_forward(
|
|
&self,
|
|
states: &Tensor,
|
|
deterministic: bool,
|
|
) -> Result<(Tensor, Tensor)> {
|
|
let batch_size = states.shape().dims()[0];
|
|
|
|
if deterministic {
|
|
// Deterministic action (mean of policy) - simplified for testing
|
|
let actions = Tensor::randn(&[batch_size, self.action_dim], &self.device)?;
|
|
let log_probs = Tensor::zeros([batch_size], &self.device)?;
|
|
Ok((actions, log_probs))
|
|
} else {
|
|
// Stochastic action with reparameterization trick - simplified for testing
|
|
let actions = Tensor::randn(&[batch_size, self.action_dim], &self.device)?;
|
|
let log_probs = Tensor::full(&[batch_size], -2.0, &self.device)?; // Dummy negative log probs
|
|
Ok((actions, log_probs))
|
|
}
|
|
}
|
|
|
|
pub async fn critic_forward(
|
|
&self,
|
|
states: &Tensor,
|
|
_actions: &Tensor,
|
|
_critic_id: usize,
|
|
) -> Result<Tensor> {
|
|
let batch_size = states.shape().dims()[0];
|
|
|
|
// Dummy Q-value computation
|
|
let q_values = Tensor::randn(&[batch_size], &self.device)?;
|
|
Ok(q_values)
|
|
}
|
|
|
|
pub async fn target_critic_forward(
|
|
&self,
|
|
states: &Tensor,
|
|
actions: &Tensor,
|
|
critic_id: usize,
|
|
) -> Result<Tensor> {
|
|
// Same as critic_forward but using target parameters
|
|
self.critic_forward(states, actions, critic_id).await
|
|
}
|
|
|
|
pub async fn compute_actor_loss(&self, states: &Tensor) -> Result<Tensor> {
|
|
let (actions, _log_probs) = self.actor_forward(states, false).await?;
|
|
|
|
let q1_values = self.critic_forward(states, &actions, 0).await?;
|
|
let _q2_values = self.critic_forward(states, &actions, 1).await?;
|
|
// Simplified for RED phase
|
|
let _min_q_values = q1_values;
|
|
|
|
let _alpha: f32 = self.log_alpha.exp()?.item()?;
|
|
// Simplified for RED phase
|
|
let actor_loss = Tensor::full(&[1], 0.1, &self.device)?;
|
|
|
|
Ok(actor_loss)
|
|
}
|
|
|
|
pub async fn compute_critic_loss(
|
|
&self,
|
|
states: &Tensor,
|
|
actions: &Tensor,
|
|
rewards: &Tensor,
|
|
next_states: &Tensor,
|
|
_dones: &Tensor,
|
|
) -> Result<(Tensor, Tensor)> {
|
|
// Current Q-values
|
|
let _q1_values = self.critic_forward(states, actions, 0).await?;
|
|
let _q2_values = self.critic_forward(states, actions, 1).await?;
|
|
|
|
// Target Q-values
|
|
let (next_actions, _next_log_probs) = self.actor_forward(next_states, false).await?;
|
|
let target_q1 = self
|
|
.target_critic_forward(next_states, &next_actions, 0)
|
|
.await?;
|
|
let _target_q2 = self
|
|
.target_critic_forward(next_states, &next_actions, 1)
|
|
.await?;
|
|
// Simplified for RED phase - just use one of the Q values
|
|
let min_target_q = target_q1;
|
|
|
|
let _alpha: f32 = self.log_alpha.exp()?.item()?;
|
|
// Simplified for RED phase
|
|
let _next_v = min_target_q;
|
|
|
|
// Simplified for RED phase
|
|
let _target_q = rewards.clone();
|
|
|
|
// Simplified for RED phase
|
|
let critic1_loss = Tensor::full(&[1], 0.1, &self.device)?;
|
|
let critic2_loss = Tensor::full(&[1], 0.1, &self.device)?;
|
|
|
|
Ok((critic1_loss, critic2_loss))
|
|
}
|
|
|
|
pub async fn compute_temperature_loss(&self, log_probs: &Tensor) -> Result<Tensor> {
|
|
if !self.config.automatic_entropy_tuning {
|
|
return Ok(Tensor::zeros([1], &self.device)?);
|
|
}
|
|
|
|
let target_entropy = self
|
|
.config
|
|
.target_entropy
|
|
.unwrap_or(-(self.action_dim as f32));
|
|
let _target_entropy_tensor =
|
|
Tensor::full(&[log_probs.shape().dims()[0]], target_entropy, &self.device)?;
|
|
// Simplified for RED phase
|
|
let temperature_loss = Tensor::full(&[1], 0.01, &self.device)?;
|
|
|
|
Ok(temperature_loss)
|
|
}
|
|
|
|
pub async fn get_target_critic_params(&self) -> Result<Vec<f32>> {
|
|
// Dummy implementation for testing
|
|
Ok(vec![1.0, 2.0, 3.0, 4.0])
|
|
}
|
|
|
|
pub async fn update(
|
|
&mut self,
|
|
states: &Tensor,
|
|
actions: &Tensor,
|
|
rewards: &Tensor,
|
|
next_states: &Tensor,
|
|
dones: &Tensor,
|
|
) -> Result<SACMetrics> {
|
|
// Compute losses
|
|
let actor_loss = self.compute_actor_loss(states).await?;
|
|
let (critic1_loss, critic2_loss) = self
|
|
.compute_critic_loss(states, actions, rewards, next_states, dones)
|
|
.await?;
|
|
|
|
let (_, log_probs) = self.actor_forward(states, false).await?;
|
|
let temperature_loss = self.compute_temperature_loss(&log_probs).await?;
|
|
|
|
// In real implementation, we'd perform gradient updates here
|
|
|
|
// Soft update target networks (dummy implementation)
|
|
// target = tau * current + (1 - tau) * target
|
|
|
|
let alpha: f32 = self.log_alpha.exp()?.item()?;
|
|
let target_entropy = self
|
|
.config
|
|
.target_entropy
|
|
.unwrap_or(-(self.action_dim as f32));
|
|
|
|
Ok(SACMetrics {
|
|
actor_loss: actor_loss.item()?,
|
|
critic1_loss: critic1_loss.item()?,
|
|
critic2_loss: critic2_loss.item()?,
|
|
temperature_loss: temperature_loss.item()?,
|
|
alpha,
|
|
target_entropy,
|
|
})
|
|
}
|
|
}
|