feature(zc): add MetaDiffuser and prompt-dt#771
Open
Super1ce wants to merge 17 commits intoopendilab:mainfrom
Open
feature(zc): add MetaDiffuser and prompt-dt#771Super1ce wants to merge 17 commits intoopendilab:mainfrom
Super1ce wants to merge 17 commits intoopendilab:mainfrom
Conversation
zjowowen
reviewed
Feb 5, 2024
| ) -> 'Policy': # noqa | ||
| """ | ||
| Overview: | ||
| Serial pipeline entry. |
zjowowen
reviewed
Feb 5, 2024
| # use the original batch size per gpu and increase learning rate | ||
| # correspondingly. | ||
| cfg.policy.learn.batch_size // get_world_size(), | ||
| # cfg.policy.learn.batch_size |
zjowowen
reviewed
Feb 5, 2024
| for epoch in range(cfg.policy.learn.train_epoch): | ||
| if get_world_size() > 1: | ||
| dataloader.sampler.set_epoch(epoch) | ||
| for i in range(cfg.policy.train_num): |
Collaborator
There was a problem hiding this comment.
"train_num"->"batch_size"?
zjowowen
reviewed
Feb 5, 2024
| (prompt_returns_embeddings, prompt_state_embeddings, prompt_action_embeddings), dim=1 | ||
| ).permute(0, 2, 1, 3).reshape(prompt_states.shape[0], 3 * prompt_seq_length, self.h_dim) | ||
|
|
||
| # prompt_stacked_attention_mask = torch.stack( |
Collaborator
There was a problem hiding this comment.
Remove these unused lines?
zjowowen
reviewed
Feb 5, 2024
ding/model/template/diffusion.py
Outdated
| self.returns_condition = returns_condition | ||
| self.condition_guidance_w = condition_guidance_w | ||
|
|
||
| # def get_loss_weights(self, discount: int): |
Collaborator
There was a problem hiding this comment.
Remove these unused lines?
zjowowen
reviewed
Feb 5, 2024
|
|
||
| return model_mean + model_std * noise, y | ||
|
|
||
| def free_guidance_sample( |
Collaborator
There was a problem hiding this comment.
Add class hints for all arguments, add Overview for functions and classes.
zjowowen
reviewed
Feb 5, 2024
ding/model/template/diffusion.py
Outdated
|
|
||
| self.embed = nn.Sequential( | ||
| nn.Linear((obs_dim * 2 + action_dim + 1) * encoder_horizon, dim * 4), | ||
| Mish(),#nn.Mish(), |
zjowowen
reviewed
Feb 5, 2024
| self._learn_model = model_wrap(self._model, wrapper_name='base') | ||
| self._learn_model.reset() | ||
|
|
||
| def _forward_learn(self, data: List[torch.Tensor]) -> Dict[str, Any]: |
Collaborator
There was a problem hiding this comment.
data should be collated into batchsize before entering policy._forward_learn.
data type shoule be Dict[str, torch.Tensor].
zjowowen
reviewed
Feb 5, 2024
| if self.have_train: | ||
| if self.task_id is None: | ||
| self.task_id = [0] * self.eval_batch_size | ||
| # if data_id is None: |
zjowowen
reviewed
Feb 5, 2024
| if self._cuda: | ||
| data = to_device(data, self._device) | ||
|
|
||
| p_s, p_a, p_rtg, p_t, p_mask, timesteps, states, actions, rewards, returns_to_go, \ |
Collaborator
There was a problem hiding this comment.
data should be collated into batchsize before entering policy._forward_learn.
data type shoule be Dict[str, torch.Tensor], so that it can be assigned confirmly.
zjowowen
reviewed
Feb 5, 2024
| self.returns_mlp = nn.Sequential( | ||
| SinusoidalPosEmb(dim), | ||
| nn.Linear(dim, dim * 4), | ||
| #nn.Mish(), |
zjowowen
reviewed
Feb 5, 2024
|
|
||
| @DATASET_REGISTRY.register('meta_traj') | ||
| class MetaTraj(Dataset): | ||
| def __init__(self, cfg): |
Collaborator
There was a problem hiding this comment.
Add notation for this class and config items.
zjowowen
reviewed
Feb 5, 2024
| Interaction serial evaluator class, policy interacts with env. This class evaluator algorithm | ||
| with test environment list. | ||
| Interfaces: | ||
| __init__, reset, reset_policy, reset_env, close, should_eval, eval |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add MetaDIffusion and prompt-dt algorithm