8000 fix(policy): saving of temperature init on SAC by AdilZouitine · Pull Request #1373 · huggingface/lerobot · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

fix(policy): saving of temperature init on SAC #1373

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
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(policy): saving of temperature init on SAC #1373

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
13 changes: 13 additions & 0 deletions lerobot/common/policies/sac/modeling_sac.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,18 @@
# limitations under the License.

import math
from copy import deepcopy
from dataclasses import asdict
from pathlib import Path
from typing import Callable, Literal

import einops
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F # noqa: N812
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
from safetensors.torch import save_model as save_model_as_safetensor
from torch import Tensor
from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution

Expand Down Expand Up @@ -480,6 +484,15 @@ def _init_temperature(self):
self.log_alpha = nn.Parameter(torch.tensor([math.log(temp_init)]))
self.temperature = self.log_alpha.exp().item()

def _save_pretrained(self, save_directory: Path) -> None:
# HACK: In order to resume training, we need to save the config the current temperature
# as initial temperature
config = deepcopy(self.config)
config.temperature_init = self.temperature
config._save_pretrained(save_directory)
model_to_save = self.module if hasattr(self, "module") else self
save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE))


class SACObservationEncoder(nn.Module):
"""Encode image and/or state vector observations."""
Expand Down
Loading
0