rewind-reward-pytorch

ReWiND Reward - Pytorch (wip)

Implementation of ReWiND, "Language-Guided Rewards Teach Robot Policies without New Demonstrations", from USC / Amazon Robotics

Install

$ pip install rewind-reward-pytorch

Usage

import torch

from rewind_reward_pytorch import (
    RewardModel,
    RewindTrainWrapper
)

reward_model = RewardModel(
    reward_bins = 10,
    categorical_rewards = False,

)

wrapper = RewindTrainWrapper(
    reward_model,
    rewind_augmentation_prob = 0.25
)

commands = [
  'pick up the blue ball and put it in the red tray',
  'pick up the red cube and put it in the green bin',
]

video = torch.rand(2, 3, 16, 224, 224)

loss = wrapper(
    commands,
    video
)

loss.backward()

# after much training

pred_rewards = reward_model(
    commands,
    video
) # (2, 16) - per frame progress prediction

assert pred_rewards.shape == (2, 16)

Citations

@article{Zhang2025ReWiNDLR,
    title   = {ReWiND: Language-Guided Rewards Teach Robot Policies without New Demonstrations},
    author  = {Jiahui Zhang and Yusen Luo and Abrar Anwar and Sumedh Anand Sontakke and Joseph J. Lim and Jesse Thomason and Erdem Biyik and Jesse Zhang},
    journal = {ArXiv},
    year    = {2025},
    volume  = {abs/2505.10911},
    url     = {https://api.semanticscholar.org/CorpusID:278714746}
}

Metadata

Release files for rewind-reward-pytorch 0.0.25

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for rewind-reward-pytorch 0.0.25
File	Size	Uploaded
rewind_reward_pytorch-0.0.25.tar.gz	910.1 kB	Aug 15, 2025	Details

Built distribution (wheel)

Table of built distributions (wheels) for rewind-reward-pytorch 0.0.25
File	Interpreter	ABI	Platform	Reset
rewind_reward_pytorch-0.0.25-py3-none-any.whl 7.3 kB Aug 15, 2025	Python 3	none	any	Details

Total release size: 917.3 kB

Release files / rewind_reward_pytorch-0.0.25.tar.gz

Download URL	rewind_reward_pytorch-0.0.25.tar.gz
Size	910.1 kB
Tags	Source
SHA-256 checksum How to use checksums	`ae695ca04f8bc035ebfda688c8b04035490482315d84d1ae1f4f2ec6aa3abb53`
BLAKE2b-256 checksum How to use checksums	`04393b7d9bd2aa813e3a7fea6568e7908bec45c5ece712bf08651e2a0135eb27`
Upload date	Aug 15, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.23`

Release files / rewind_reward_pytorch-0.0.25-py3-none-any.whl

Download URL	rewind_reward_pytorch-0.0.25-py3-none-any.whl
Size	7.3 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`30431f90afbba37442591f1205bcb2d18e6bb5990f4c9099ec25adc9baf1beec`
BLAKE2b-256 checksum How to use checksums	`208cd003cc0ecea83f4fc628466cb9cded8195e85ea7403350afed45a16dcf9a`
Upload date	Aug 15, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.23`