57 lines
1.4 KiB
Python
57 lines
1.4 KiB
Python
import torch
|
|
import logging
|
|
from model.common.gmm import GMMModel
|
|
|
|
|
|
class VPG_GMM(GMMModel):
|
|
def __init__(
|
|
self,
|
|
actor,
|
|
critic,
|
|
cond_steps=1,
|
|
network_path=None,
|
|
**kwargs,
|
|
):
|
|
super().__init__(network=actor, **kwargs)
|
|
self.cond_steps = cond_steps
|
|
|
|
# Re-name network to actor
|
|
self.actor_ft = actor
|
|
|
|
# Value function for obs - simple MLP
|
|
self.critic = critic.to(self.device)
|
|
if network_path is not None:
|
|
checkpoint = torch.load(
|
|
network_path, map_location=self.device, weights_only=True
|
|
)
|
|
self.load_state_dict(
|
|
checkpoint["model"],
|
|
strict=False,
|
|
)
|
|
logging.info("Loaded actor from %s", network_path)
|
|
|
|
def get_logprobs(
|
|
self,
|
|
cond,
|
|
actions,
|
|
):
|
|
B, T, D = actions.shape
|
|
dist, entropy, std = self.forward_train(
|
|
cond.view(B, -1),
|
|
deterministic=False,
|
|
)
|
|
log_prob = dist.log_prob(actions.view(B, -1))
|
|
return log_prob, entropy, std
|
|
|
|
def loss(self, obs, chains, reward):
|
|
raise NotImplementedError
|
|
|
|
# override to diffuse over action only
|
|
@torch.no_grad()
|
|
def forward(self, cond, deterministic=False):
|
|
B = cond.shape[0]
|
|
return super().forward(
|
|
cond=cond.view(B, -1),
|
|
deterministic=deterministic,
|
|
)
|