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

View File

@@ -23,6 +23,7 @@ from typing import Any
import gymnasium as gym
import numpy as np
import torch
from gymnasium import spaces
from lerobot.types import RobotObservation
@@ -320,7 +321,11 @@ class RoboTwinEnv(gym.Env):
actual_seed = self.episode_index if seed is None else seed
setup_kwargs = _load_robotwin_setup_kwargs(self.task_name)
setup_kwargs.update(seed=actual_seed, is_test=True)
self._env.setup_demo(**setup_kwargs)
# 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.episode_index += self._reset_stride
self._step_count = 0
@@ -333,10 +338,13 @@ class RoboTwinEnv(gym.Env):
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.
if hasattr(self._env, "take_action"):
self._env.take_action(action)
else:
self._env.step(action)
# 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"):
self._env.take_action(action)
else:
self._env.step(action)
self._step_count += 1