@torch.no_grad()
def predict_V(self, x: torch.Tensor) -> np.ndarray:
    self.eval()
    x = self.trunk(x)
    V_hat = self.value_head(x)
    self.train()
    return V_hat.detach().cpu().numpy()
@torch.no_grad()
def predict_V(self, x: torch.Tensor) -> np.ndarray:
    self.eval()
    x = self.trunk(x)
    V_hat = self.value_head(x)
    self.train()
    return V_hat.detach().cpu().numpy()