Skip to content

Warn if using offline data buffer without Agent wrapper #1771

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

Open
QuantuMope wants to merge 3 commits into
base: pytorch
Choose a base branch
from
Open

Warn if using offline data buffer without Agent wrapper #1771

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
74240c4
[Bug Fix] Enable hybrid SAC training
QuantuMope May 12, 2025
2c7e7cf
address comments
QuantuMope May 12, 2025
9db5652
address comments
QuantuMope May 13, 2025
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
4 changes: 3 additions & 1 deletion alf/algorithms/algorithm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from torch.nn.modules.module import _IncompatibleKeys, _addindent

import alf
from alf.data_structures import AlgStep, LossInfo, StepType, TimeStep
from alf.data_structures import AlgStep, LossInfo, StepType, TimeStep, BasicRolloutInfo
from alf.experience_replayers.replay_buffer import BatchInfo, ReplayBuffer
from alf.optimizers.utils import GradientNoiseScaleEstimator
from alf.utils.checkpoint_utils import (is_checkpoint_enabled,
Expand Down Expand Up @@ -1368,6 +1368,8 @@ def train_step_offline(self, inputs, state, rollout_info, pre_train=False):
customized training.
"""
try:
if isinstance(rollout_info, BasicRolloutInfo):
rollout_info = rollout_info.rl
return self.train_step(inputs, state, rollout_info)
except:
# the default train_step is not compatible with the
Expand Down
14 changes: 8 additions & 6 deletions alf/algorithms/sac_algorithm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,13 @@
from alf.algorithms.config import TrainerConfig
from alf.algorithms.off_policy_algorithm import OffPolicyAlgorithm
from alf.algorithms.one_step_loss import OneStepTDLoss
from alf.algorithms.rl_algorithm import RLAlgorithm
from alf.data_structures import TimeStep, Experience, LossInfo, namedtuple
from alf.data_structures import TimeStep, LossInfo, namedtuple, \
BasicRLInfo
from alf.data_structures import AlgStep, StepType
from alf.nest import nest
import alf.nest.utils as nest_utils
from alf.networks import ActorDistributionNetwork, CriticNetwork
from alf.networks import QNetwork, QRNNNetwork
from alf.networks import QNetwork
from alf.tensor_specs import TensorSpec, BoundedTensorSpec
from alf.utils import losses, common, dist_utils, math_ops
from alf.utils.normalizers import ScalarAdaptiveNormalizer
Expand Down Expand Up @@ -845,8 +845,10 @@ def _select_q_value(self, action, q_values):
return q_values.gather(2, action).squeeze(2)

def _critic_train_step(self, observation, target_observation,
state: SacCriticState, rollout_info: SacInfo,
action, action_distribution):
state: SacCriticState,
rollout_info: SacInfo | BasicRLInfo, action,

@hnyu hnyu May 12, 2025 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still should always be SacInfo? If it's BasicRLInfo, the algorithm will crash.

@QuantuMope QuantuMope May 12, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Offline buffer data is stored as BasicRLInfo which comprises of just (s,a,r) data.

@QuantuMope QuantuMope May 12, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess, I could convert BasicRLInfo into SacInfo with some fields empty? Not sure which is a better design. Lmk which one you think is cleaner and I can change.

@hnyu hnyu May 12, 2025 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Offline buffer data is stored as BasicRLInfo which comprises of just (s,a,r) data.

If you look at SAC's train_step(), it will get access to rollout_info.repr. This means that SAC is currently incompatible with offline training.

@QuantuMope QuantuMope May 12, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm using a frozen encoder so not training a repr.

@QuantuMope QuantuMope May 12, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wait, it seems that repr is stored in BasicRolloutInfo. Not sure how this code was running then. I'll take a look

@hnyu hnyu May 12, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's using elastic_namedtuple so any missing field returns (). Anyway, a little weird but it works.

@Haichao-Zhang Haichao-Zhang May 12, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally, we should not include BasicRLInfo here as it could confuse the pure sac users. The better alternative might be comply with the Agent assumption and possibly extend it.

@QuantuMope QuantuMope May 13, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed typehints. Also added a warning message advising users to use Agent. Before the code would simply crash due to interface conflict.

action_distribution):

critics, critics_state = self._compute_critics(
self._critic_networks,
observation,
Expand Down Expand Up @@ -897,7 +899,7 @@ def _alpha_train_step(self, log_pi):
return sum(nest.flatten(alpha_loss))

def train_step(self, inputs: TimeStep, state: SacState,
rollout_info: SacInfo):
rollout_info: SacInfo | BasicRLInfo):
assert not self._is_eval
self._training_started = True
if self._target_repr_alg is not None:
Expand Down