Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BugFix, Feature] Vmap randomness in losses #1740

Merged
merged 21 commits into from
Jan 9, 2024
Merged

[BugFix, Feature] Vmap randomness in losses #1740

Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
3206ddb
add randomness param to vmap for offpolicy algos
BY571 Dec 7, 2023
82ec897
update td3 with vmap_randomness
BY571 Dec 8, 2023
29a5cda
undo vmap randomness zip actor and critic modules
BY571 Dec 11, 2023
53ca415
expand random module list
BY571 Dec 11, 2023
50aa8a4
move random_module_list to utils
BY571 Dec 11, 2023
3f54635
update sac losses
BY571 Dec 13, 2023
c4ff926
update iql objective
BY571 Dec 13, 2023
1c44c35
update redq objective
BY571 Dec 13, 2023
bc03b92
update conti cql
BY571 Dec 13, 2023
9aa65ed
move vmap_randomness to LossModule
BY571 Dec 14, 2023
3ff3ef8
fix
BY571 Dec 14, 2023
febf277
fix
BY571 Dec 14, 2023
3af72c5
add test example
BY571 Dec 14, 2023
08c83c1
Merge branch 'main' into vmap_dropout
BY571 Dec 27, 2023
f122825
fix
BY571 Dec 27, 2023
d85d63b
add vmap randomness test
BY571 Jan 3, 2024
92c3e40
update ranodm_module_list
BY571 Jan 4, 2024
7cdc1e6
add fail case for vmap
BY571 Jan 4, 2024
f737172
Merge branch 'main' into vmap_dropout
BY571 Jan 8, 2024
1bfcc27
update vmap fail case test
BY571 Jan 8, 2024
6ac65e1
Merge remote-tracking branch 'origin/main' into vmap_dropout
vmoens Jan 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
update redq objective
  • Loading branch information
@BY571
BY571 committed Dec 13, 2023
commit 1c44c35c3665fce236e792cad9eeaa0352a53b57
31 changes: 28 additions & 3 deletions torchrl/objectives/redq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
_vmap_func,
default_value_kwargs,
distance_loss,
RANDOM_MODULE_LIST,
ValueEstimators,
)
from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
Expand Down Expand Up @@ -234,6 +235,7 @@ class _AcceptedKeys:
"next.state_value",
"target_value",
]
_vmap_randomness = None
BY571 marked this conversation as resolved.
Show resolved Hide resolved

def __init__(
self,
Expand All @@ -255,7 +257,6 @@ def __init__(
priority_key: str = None,
separate_losses: bool = False,
):

super().__init__()
self._in_keys = None
self._set_deprecated_ctor_keys(priority_key=priority_key)
Expand Down Expand Up @@ -319,9 +320,11 @@ def __init__(
self.gamma = gamma

self._vmap_qvalue_network00 = _vmap_func(
self.qvalue_network, randomness="different"
self.qvalue_network, randomness=self.vmap_randomness
)
self._vmap_getdist = _vmap_func(
self.actor_network, func="get_dist_params", randomess=self.vmap_randomness
)
self._vmap_getdist = _vmap_func(self.actor_network, func="get_dist_params")

@property
def target_entropy(self):
Expand Down Expand Up @@ -406,6 +409,28 @@ def in_keys(self):
def in_keys(self, values):
self._in_keys = values

@property
def vmap_randomness(self):
if self._vmap_randomness is None:
do_break = False
for val in self.__dict__.values():
if isinstance(val, torch.nn.Module):
for module in val.modules():
if isinstance(module, RANDOM_MODULE_LIST):
self._vmap_randomness = "different"
do_break = True
break
if do_break:
# double break
break
else:
self._vmap_randomness = "error"

return self._vmap_randomness

def set_vmap_randomness(self, value):
self._vmap_randomness = value

BY571 marked this conversation as resolved.
Show resolved Hide resolved
@property
@_cache_values
def _cached_detach_qvalue_network_params(self):
Expand Down
Loading