Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.distributions import Normal | |
| import torch.optim as optim | |
| from .rssm import RSSMState, RSSM | |
| from .networks import Encoder, Decoder, RewardModel, Actor, Critic | |
| def compute_lambda_returns(rewards: torch.Tensor, values: torch.Tensor, gamma: float = 0.99, lambda_: float = 0.95): | |
| """ | |
| Computes TD(lambda) returns over an imagined trajectory. | |
| rewards: [horizon, batch] | |
| values: [horizon, batch] | |
| """ | |
| horizon = rewards.shape[0] | |
| returns = [] | |
| last_return = values[-1] | |
| for t in reversed(range(horizon - 1)): | |
| # TD(lambda) recursive calculation | |
| ret = rewards[t] + gamma * ((1 - lambda_) * values[t+1] + lambda_ * last_return) | |
| returns.insert(0, ret) | |
| last_return = ret | |
| returns.append(values[-1]) | |
| return torch.stack(returns) | |
| class AgenticForecaster(nn.Module): | |
| def __init__(self, obs_dim: int, action_dim: int, device: torch.device): | |
| super().__init__() | |
| self.device = device | |
| self.action_dim = action_dim | |
| self.encoder = Encoder(obs_dim=obs_dim).to(device) | |
| self.rssm = RSSM(action_dim=action_dim).to(device) | |
| feature_dim = self.rssm.stoch_dim + self.rssm.deter_dim | |
| self.decoder = Decoder(feature_dim=feature_dim, obs_dim=obs_dim).to(device) | |
| self.reward_model = RewardModel(feature_dim=feature_dim).to(device) | |
| self.actor = Actor(feature_dim=feature_dim, action_dim=action_dim).to(device) | |
| self.critic = Critic(feature_dim=feature_dim).to(device) | |
| self.target_critic = Critic(feature_dim=feature_dim).to(device) | |
| self.target_critic.load_state_dict(self.critic.state_dict()) | |
| # Optimizers | |
| self.model_optimizer = optim.Adam( | |
| list(self.encoder.parameters()) + | |
| list(self.rssm.parameters()) + | |
| list(self.decoder.parameters()) + | |
| list(self.reward_model.parameters()), | |
| lr=3e-4 | |
| ) | |
| self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=1e-4) | |
| self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=1e-4) | |
| def train_world_model(self, observations: torch.Tensor, actions: torch.Tensor, rewards: torch.Tensor, free_nats: float = 3.0, kl_scale: float = 1.0): | |
| """ | |
| observations: [batch_size, seq_len, obs_dim] | |
| actions: [batch_size, seq_len, action_dim] | |
| rewards: [batch_size, seq_len] | |
| """ | |
| batch_size, seq_len, _ = observations.shape | |
| state = self.rssm.initial_state(batch_size, self.device) | |
| kl_losses = [] | |
| recon_losses = [] | |
| reward_losses = [] | |
| for t in range(seq_len): | |
| obs_embed = self.encoder(observations[:, t, :]) | |
| # Actions align such that action[t-1] led to observation[t]. For t=0, assume zero action | |
| action_t = actions[:, t-1, :] if t > 0 else torch.zeros(batch_size, self.action_dim, device=self.device) | |
| state, post_stats, prior_stats = self.rssm.observe_step(state, action_t, obs_embed) | |
| features = state.get_features() | |
| recon_obs = self.decoder(features) | |
| pred_reward = self.reward_model(features) | |
| recon_losses.append(F.mse_loss(recon_obs, observations[:, t, :], reduction='none').sum(-1).mean()) | |
| reward_losses.append(F.mse_loss(pred_reward.squeeze(-1), rewards[:, t], reduction='mean')) | |
| post_dist = Normal(post_stats[0], post_stats[1]) | |
| prior_dist = Normal(prior_stats[0], prior_stats[1]) | |
| # Analytical KL divergence summed over latent dimension, averaged over batch | |
| kl_loss = torch.mean(torch.sum(torch.distributions.kl.kl_divergence(post_dist, prior_dist), dim=-1)) | |
| # Free bits (nats) implementation | |
| kl_loss = torch.max(kl_loss, torch.tensor(free_nats).to(self.device)) | |
| kl_losses.append(kl_loss) | |
| model_loss = sum(recon_losses) + sum(reward_losses) + kl_scale * sum(kl_losses) | |
| self.model_optimizer.zero_grad() | |
| model_loss.backward() | |
| nn.utils.clip_grad_norm_(self.parameters(), 100.0) | |
| self.model_optimizer.step() | |
| return {"model_loss": model_loss.item(), "recon_loss": sum(recon_losses).item(), "kl_loss": sum(kl_losses).item()} | |
| def train_actor_critic(self, start_states: RSSMState, horizon: int = 15, gamma: float = 0.99, lambda_: float = 0.95): | |
| """ | |
| Trains Actor and Critic purely in latent imagination. | |
| start_states is sampled from the posterior distribution during world model training. | |
| """ | |
| # 1. Rollout imagined trajectory | |
| state = start_states | |
| features_list = [] | |
| rewards_list = [] | |
| for _ in range(horizon): | |
| features = state.get_features() | |
| action = self.actor(features.detach(), explore=True) | |
| # Imagine next state driven by the actor's action | |
| state = self.rssm.imagine_step(state, action) | |
| next_features = state.get_features() | |
| reward = self.reward_model(next_features) | |
| features_list.append(next_features) | |
| rewards_list.append(reward.squeeze(-1)) | |
| features_stack = torch.stack(features_list) # [horizon, batch, feature_dim] | |
| rewards_stack = torch.stack(rewards_list) # [horizon, batch] | |
| # 2. Compute TD(lambda) values using the target critic | |
| with torch.no_grad(): | |
| values_stack = self.target_critic(features_stack).squeeze(-1) | |
| returns = compute_lambda_returns(rewards_stack, values_stack, gamma, lambda_) | |
| # 3. Train Actor (Maximize expected lambda return) | |
| actor_loss = -returns.mean() | |
| self.actor_optimizer.zero_grad() | |
| actor_loss.backward() | |
| nn.utils.clip_grad_norm_(self.actor.parameters(), 100.0) | |
| self.actor_optimizer.step() | |
| # 4. Train Critic (Predict lambda returns) | |
| # Detach features so we don't backprop into the world model | |
| pred_values = self.critic(features_stack.detach()).squeeze(-1) | |
| critic_loss = F.mse_loss(pred_values, returns.detach()) | |
| self.critic_optimizer.zero_grad() | |
| critic_loss.backward() | |
| nn.utils.clip_grad_norm_(self.critic.parameters(), 100.0) | |
| self.critic_optimizer.step() | |
| # 5. Soft update target critic | |
| tau = 0.05 | |
| for param, target_param in zip(self.critic.parameters(), self.target_critic.parameters()): | |
| target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data) | |
| return {"actor_loss": actor_loss.item(), "critic_loss": critic_loss.item()} | |
| def act(self, obs: torch.Tensor, prev_state: RSSMState, prev_action: torch.Tensor, explore: bool = False): | |
| """Used during data collection or evaluation to select an action.""" | |
| with torch.no_grad(): | |
| obs_embed = self.encoder(obs.unsqueeze(0)) | |
| new_state, _, _ = self.rssm.observe_step(prev_state, prev_action.unsqueeze(0), obs_embed) | |
| features = new_state.get_features() | |
| action = self.actor(features, explore=explore) | |
| return action.squeeze(0), new_state | |
| class HomeostaticAgenticForecaster(AgenticForecaster): | |
| """ | |
| Dreamer agent with homeostatic critic. | |
| Maintains a stable value baseline while learning deviations. | |
| """ | |
| def __init__(self, obs_dim: int, action_dim: int, device: torch.device): | |
| super().__init__(obs_dim, action_dim, device) | |
| # Replace critic with homeostatic version | |
| feature_dim = self.rssm.stoch_dim + self.rssm.deter_dim | |
| from .networks import HomeostaticCritic | |
| self.critic = HomeostaticCritic(feature_dim=feature_dim).to(device) | |
| self.target_critic = HomeostaticCritic(feature_dim=feature_dim).to(device) | |
| self.target_critic.load_state_dict(self.critic.state_dict()) | |
| # Reinitialize optimizers with new parameters | |
| self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=1e-4) | |
| self.training_step = 0 | |
| def train_actor_critic(self, start_states: RSSMState, horizon: int = 15, | |
| gamma: float = 0.99, lambda_: float = 0.95): | |
| """Same as base, but critic updates its setpoint periodically""" | |
| # After computing returns, update homeostasis | |
| if self.training_step % 100 == 0: # Every 100 steps | |
| update_homeostasis = True | |
| else: | |
| update_homeostasis = False | |
| self.training_step += 1 | |
| # 1. Rollout imagined trajectory | |
| state = start_states | |
| features_list = [] | |
| rewards_list = [] | |
| for _ in range(horizon): | |
| features = state.get_features() | |
| action = self.actor(features.detach(), explore=True) | |
| state = self.rssm.imagine_step(state, action) | |
| next_features = state.get_features() | |
| reward = self.reward_model(next_features) | |
| features_list.append(next_features) | |
| rewards_list.append(reward.squeeze(-1)) | |
| features_stack = torch.stack(features_list) | |
| rewards_stack = torch.stack(rewards_list) | |
| # 2. Compute TD(lambda) values using the target critic | |
| with torch.no_grad(): | |
| values_stack = self.target_critic(features_stack).squeeze(-1) | |
| returns = compute_lambda_returns(rewards_stack, values_stack, gamma, lambda_) | |
| # 3. Train Actor | |
| actor_loss = -returns.mean() | |
| self.actor_optimizer.zero_grad() | |
| actor_loss.backward() | |
| nn.utils.clip_grad_norm_(self.actor.parameters(), 100.0) | |
| self.actor_optimizer.step() | |
| # 4. Train Critic (Predict lambda returns AND update setpoint optionally) | |
| pred_values = self.critic(features_stack.detach(), update_homeostasis=update_homeostasis).squeeze(-1) | |
| critic_loss = F.mse_loss(pred_values, returns.detach()) | |
| self.critic_optimizer.zero_grad() | |
| critic_loss.backward() | |
| nn.utils.clip_grad_norm_(self.critic.parameters(), 100.0) | |
| self.critic_optimizer.step() | |
| # 5. Soft update target critic | |
| tau = 0.05 | |
| for param, target_param in zip(self.critic.parameters(), self.target_critic.parameters()): | |
| target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data) | |
| return {"actor_loss": actor_loss.item(), "critic_loss": critic_loss.item()} | |