fix(robotwin): re-enable autograd for CuRobo planner warmup and take_action

lerobot_eval wraps the full rollout in torch.no_grad() (lerobot_eval.py:566),
but RoboTwin's setup_demo → load_robot → CuroboPlanner(...) runs
motion_gen.warmup(), which invokes Newton's-method trajectory optimization.
That optimizer calls cost.backward() internally, which raises

    RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

when autograd is disabled. take_action() hits the same planner path at every
step. Wrap both setup_demo and take_action in torch.enable_grad() so CuRobo's
optimizer can build its computation graph. Policy inference is unaffected —
rollout()'s inner torch.inference_mode() block around select_action() is
untouched, so we still don't allocate grad buffers during policy forward.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-04-14 17:39:21 +02:00
parent 5558ea2207
commit b06f134fe4
+8
View File
@@ -23,6 +23,7 @@ from typing import Any
import gymnasium as gym import gymnasium as gym
import numpy as np import numpy as np
import torch
from gymnasium import spaces from gymnasium import spaces
from lerobot.types import RobotObservation from lerobot.types import RobotObservation
@@ -320,6 +321,10 @@ class RoboTwinEnv(gym.Env):
actual_seed = self.episode_index if seed is None else seed actual_seed = self.episode_index if seed is None else seed
setup_kwargs = _load_robotwin_setup_kwargs(self.task_name) setup_kwargs = _load_robotwin_setup_kwargs(self.task_name)
setup_kwargs.update(seed=actual_seed, is_test=True) setup_kwargs.update(seed=actual_seed, is_test=True)
# setup_demo → load_robot → CuroboPlanner.warmup() runs Newton's-method
# trajectory optimization, which requires autograd. lerobot_eval wraps
# the whole rollout in torch.no_grad(), so re-enable grad here.
with torch.enable_grad():
self._env.setup_demo(**setup_kwargs) self._env.setup_demo(**setup_kwargs)
self.episode_index += self._reset_stride self.episode_index += self._reset_stride
self._step_count = 0 self._step_count = 0
@@ -333,6 +338,9 @@ class RoboTwinEnv(gym.Env):
raise ValueError(f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}") raise ValueError(f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}")
# RoboTwin 2.0 uses take_action(); fall back to step() for older forks. # RoboTwin 2.0 uses take_action(); fall back to step() for older forks.
# take_action() invokes the CuRobo planner, which needs autograd —
# lerobot_eval wraps the rollout in torch.no_grad().
with torch.enable_grad():
if hasattr(self._env, "take_action"): if hasattr(self._env, "take_action"):
self._env.take_action(action) self._env.take_action(action)
else: else: