From 188bd201aa2ce5e7b0cfc224d470935311a1a1d1 Mon Sep 17 00:00:00 2001
From: Nicklas Hansen <hello@nicklashansen.com>
Date: Thu, 4 Jan 2024 19:39:44 -0800
Subject: [PATCH] disable uncertainty estimation when coef=0

---
 tdmpc2/tdmpc2.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tdmpc2/tdmpc2.py b/tdmpc2/tdmpc2.py
index c70ee2e..797c58d 100755
--- a/tdmpc2/tdmpc2.py
+++ b/tdmpc2/tdmpc2.py
@@ -94,6 +94,8 @@ class TDMPC2:
 	@torch.no_grad()
 	def _estimate_uncertainty(self, z, task):
 		"""Estimates epistemic uncertainty, normalized by predicted value."""
+		if self.cfg.uncertainty_coef == 0:
+			return 0
 		qs = math.two_hot_inv(self.model.Q(z, self.model.pi(z, task)[1], task, return_type='all'), self.cfg)
 		return qs.mean() * qs.std(0) * self.cfg.uncertainty_coef