Skip to content

Commit

Permalink
polish(mark): add hybrid action space support to ActionNoiseWrapper (#…
Browse files Browse the repository at this point in the history
…829)

* Update model_wrappers.py

Add hybrid action space support to action noise wrapper.

* Updated syntax and added comment
  • Loading branch information
MarkHolmstrom authored Sep 20, 2024
1 parent 44d5479 commit 6ae1396
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions ding/model/wrapper/model_wrappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -866,10 +866,14 @@ def forward(self, *args, **kwargs):
assert isinstance(output, dict), "model output must be dict, but find {}".format(type(output))
if 'action' in output or 'action_args' in output:
key = 'action' if 'action' in output else 'action_args'
action = output[key]
# handle hybrid action space by adding noise to continuous part of model output
action = output[key]['action_args'] if isinstance(output[key], dict) else output[key]
assert isinstance(action, torch.Tensor)
action = self.add_noise(action)
output[key] = action
if isinstance(output[key], dict):
output[key]['action_args'] = action
else:
output[key] = action
return output

def add_noise(self, action: torch.Tensor) -> torch.Tensor:
Expand Down

0 comments on commit 6ae1396

Please sign in to comment.