ydy0615/lerobot-clone
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-06-02 20:01:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

hl-gauss

Browse Source
This commit is contained in:
Pepijn
2025-09-01 13:34:55 +02:00
parent 45348d7b69
commit a305f5f46a
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/lerobot/policies/rlearn/modeling_rlearn.py
View File

@@ -422,6 +422,7 @@ class RLearNPolicy(PreTrainedPolicy):
# Compute main loss (or just return predictions in eval)
loss_start = time.perf_counter()
loss = torch.tensor(0.0, device=device)
if target is None:
total_loss = torch.tensor(0.0, device=device)
loss = total_loss
@@ -435,6 +436,7 @@ class RLearNPolicy(PreTrainedPolicy):
bin_idx = torch.where(video_mask, bin_idx, torch.full_like(bin_idx, -1))
loss_ce = F.cross_entropy(video_frame_logits.permute(0, 2, 1), bin_idx, ignore_index=-1)
total_loss = loss_ce
loss = loss_ce
predicted_rewards = torch.softmax(video_frame_logits, dim=-1)
else:
# HL-Gauss or regression