chore: import upstream snapshot with attribution

2026-07-13 12:36:51 +08:00
commit d731240295
27 changed files with 3078 additions and 0 deletions
@@ -0,0 +1,77 @@
+# Copyright 2022 Twitter, Inc and Zhendong Wang.
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from utils.logger import logger
+
+from agents.diffusion import Diffusion
+from agents.model import MLP
+
+
+class Diffusion_BC(object):
+    def __init__(self,
+                 state_dim,
+                 action_dim,
+                 max_action,
+                 device,
+                 discount,
+                 tau,
+                 beta_schedule='linear',
+                 n_timesteps=100,
+                 lr=2e-4,
+                 ):
+
+        self.model = MLP(state_dim=state_dim, action_dim=action_dim, device=device)
+        self.actor = Diffusion(state_dim=state_dim, action_dim=action_dim, model=self.model, max_action=max_action,
+                               beta_schedule=beta_schedule, n_timesteps=n_timesteps,
+                               ).to(device)
+        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr)
+
+        self.max_action = max_action
+        self.action_dim = action_dim
+        self.discount = discount
+        self.tau = tau
+        self.device = device
+
+    def train(self, replay_buffer, iterations, batch_size=100, log_writer=None):
+
+        metric = {'bc_loss': [], 'ql_loss': [], 'actor_loss': [], 'critic_loss': []}
+        for _ in range(iterations):
+            # Sample replay buffer / batch
+            state, action, next_state, reward, not_done = replay_buffer.sample(batch_size)
+
+            loss = self.actor.loss(action, state)
+
+            self.actor_optimizer.zero_grad()
+            loss.backward()
+            self.actor_optimizer.step()
+
+            metric['actor_loss'].append(0.)
+            metric['bc_loss'].append(loss.item())
+            metric['ql_loss'].append(0.)
+            metric['critic_loss'].append(0.)
+
+        return metric
+
+    def sample_action(self, state):
+        state = torch.FloatTensor(state.reshape(1, -1)).to(self.device)
+        with torch.no_grad():
+            action = self.actor.sample(state)
+        return action.cpu().data.numpy().flatten()
+
+    def save_model(self, dir, id=None):
+        if id is not None:
+            torch.save(self.actor.state_dict(), f'{dir}/actor_{id}.pth')
+        else:
+            torch.save(self.actor.state_dict(), f'{dir}/actor.pth')
+
+    def load_model(self, dir, id=None):
+        if id is not None:
+            self.actor.load_state_dict(torch.load(f'{dir}/actor_{id}.pth'))
+        else:
+            self.actor.load_state_dict(torch.load(f'{dir}/actor.pth'))
+