FLUX 3 액션

FLUX 3 Action은 카메라 관측, 상태 및 작업 지시에서 로봇의 행동을 예측합니다. 동결된 Video VAE와 텍스트 인코더가 비디오 잠재 변수와 행동 청크를 예측하는 diffusion transformer에 입력을 제공하며, 추론 시에는 행동만 반환됩니다. LeRobot은 추론과 PEFT를 지원합니다. LoRA는 트렁크를 적응시키고, 로봇의 행동 및 조건화 헤드는 전체적으로 학습됩니다.

flux-3-action-base를 기본 액션 모델로 제공하며, SO-101 flux-3-action-so101 및 Franka Panda 로봇 flux-3-action-droid (DROID 데이터셋으로 파인튜닝됨)에 대한 파인튜닝 모델도 함께 제공합니다. 여기서 핵심 예제는 flux-3-action-so101 모델을 기반으로 LoRA를 사용하여 작업을 적응하는 것입니다.

기본 모델에서 새로운 로봇 embodiment를 학습하거나 전체 fine-tuning을 수행하려면, 독립형 flux-action 저장소를 사용하세요.

설치

LeRobot 체크아웃에서 학습, 정책, PEFT 및 EMA 의존성을 설치합니다:

pip install -e ".[training,flux3,peft,diffusion]"

Video VAE는 또한 https://whl.natten.org에서 torch/CUDA 호환 NATTEN 휠이 필요합니다. torch 2.11과 CUDA 12.8의 경우:

uv pip install natten==0.21.6+torch2110cu128 --find-links https://whl.natten.org

torch 2.10의 경우 natten==0.21.6+torch2100cu128를 사용하거나, 해당 환경에 맞게 빌드하세요.

SO-101 작업 적응

30 Hz의 LeRobot v3 데이터셋을 사용하며, 6개의 절대 명령 동작 채널과 6개의 측정 상태 채널을 동일한 관절 순서와 단위로 구성합니다. 그립퍼는 마지막에 위치합니다. 원시 값, 일관된 로봇 보정 및 지시사항을 task에 제공합니다.

카메라

카메라 이름, 개수 및 레이아웃은 설정할 수 있습니다. SO-101 패키지는 두 카메라가 촬영하는 대상에 따라 이름을 붙였습니다:

--rename_map를 사용하여 카메라 이름을 예상된 키에 매핑할 수 있습니다. 예를 들어:

--rename_map='{"observation.images.top":"observation.images.scene","observation.images.gripper_cam":"observation.images.wrist"}'

rollout에도 동일한 매핑을 사용하세요. 일치하는 키에는 매핑이 필요하지 않습니다. rename maps를 참조하세요.

로드 및 학습

black-forest-labs/flux-3-action-so101에서 SO-101 정책을 로드합니다. 저장된 구성은 black-forest-labs/flux-3-action-base의 고정된 리비전으로부터 공유 VAE와 텍스트 인코더를 자동으로 로드합니다. 액세스가 제한된 경우, 액세스 권한이 있는 계정으로 hf auth login을 사용하여 인증합니다. LeRobot 체크아웃에서 실행합니다:

lerobot-train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-so101 --policy.device=cuda \
  --dataset.repo_id=YOUR_ORG/YOUR_SO101_DATASET \
  --output_dir=outputs/so101_lora

Hub 캐시는 재개 및 어댑터 추론을 위해 다운로드된 모델과 공유 인코더를 유지합니다.

고정된 SO-101-finetuned 체크포인트는 scene / wrist 카메라 키, 8개의 관찰 타임스텝, 독립적으로 인코딩된 두 개의 시각 스냅샷, 그리고 과거 액션/상태 조건화를 사용합니다. 42개의 액션을 예측하고, 처음 32개를 30 Hz로 실행한 후 다시 계획합니다. 저장된 프로세서에는 베이스와 일치하는 정규화 통계가 포함됩니다. 이러한 체크포인트 설정은 작업 적응 과정에서 유지됩니다. 공유되는 examples/flux3/lora.json은 학습/런타임 설정과 학습 예산만 제공합니다. 동일한 파일은 각 Hub 패키지에 lora.json으로 번들링됩니다. 카메라 레이아웃, 조건화, 액션 표현 및 손실 설정은 체크포인트에서 가져옵니다. 이 명령들은 저장된 실행 길이 32, 가이던스 3/3 및 eager 실행을 유지합니다. 가이던스 4/1 및 컴파일을 사용하려면 --policy.guidance_scale=4 --policy.guidance_scale_action=1 --policy.compile_model=true를 추가하세요.

조건부 실험: LeRobot은 n_obs_steps, history_snapshots, condition_on_past_actions를 제공합니다. 이 값들을 변경하려면 호환되는 모델/프로세서 설정과 파인튜닝이 필요하며, 일반적인 작업 적응 시에는 체크포인트에 저장된 값을 유지하세요.

이 프리셋은 rank/alpha 32 LoRA와 decay 0.999의 EMA를 활성화하며, 자동으로 업로드되지 않습니다. 저장된 어댑터를 식별하는 --policy.use_peft=true가 아닌 최상위 peft 설정을 통해 새로운 어댑터를 활성화하세요.

실험에서 EMA는 DROID에서 상당한 이점을 제공한 반면, SO-101의 결과는 아직까지 혼재되어 있습니다. 이 이점은 데이터와 학습 설정에 따라 달라집니다. 배포할 모델을 선택하기 전에 EMA를 활성화한 상태로 유지하고, 동일한 rollout 프로토콜 하에서 raw 체크포인트와 EMA 체크포인트를 모두 평가하는 것을 권장합니다. 주기적인 오프라인 검증 손실은 raw 가중치를 사용하며, 두 변형을 자동으로 비교하지 않습니다.

빠른 반복을 위해 단일 GPU 예산으로 시작합니다. 더 큰 학습 예산을 사용하려면 동일한 config에 다음 overrides를 적용하세요:

accelerate launch --num_processes=4 --multi_gpu -m lerobot.scripts.lerobot_train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-so101 --policy.device=cuda \
  --dataset.repo_id=YOUR_ORG/YOUR_SO101_DATASET \
  --batch_size=8 --steps=14000 \
  --output_dir=outputs/so101_lora_large

Effective batch는 GPU 수 × GPU당 배치 크기 × accumulation입니다. LeRobot은 steps, save_freq, eval_steps를 microsteps로 계산합니다. optimizer 업데이트 횟수를 구하려면 accumulation으로 나누어야 합니다. 프리셋은 500 microsteps(125 업데이트)마다 저장 및 평가합니다. accumulation을 변경할 때는 세 설정 값 모두 accumulation으로 나누어떨어지도록 유지하세요. Scheduler의 freeze, warmup, decay, cooldown 길이도 microsteps를 사용합니다. 이러한 단계의 길이를 설정할 때는 원하는 optimizer 업데이트 횟수에 accumulation factor를 곱하세요. 두 예산은 동일한 LoRA, optimizer, EMA 설정을 사용하며, 로드된 base의 conditioning, action 차원, 정규화를 보존합니다. 체크포인트에는 raw 어댑터와 EMA 어댑터가 모두 포함됩니다. held-out 평가와 rollouts를 사용하여 예산과 체크포인트를 선택하세요. 더 큰 배치가 더 나은 결과를 보장하지는 않습니다.

프리셋은 각 작업의 마지막 20% 에피소드를 평가용으로 제외합니다. 하드웨어에 맞게 배치 크기와 축적을 조정하세요. 멀티 GPU 학습 및 PEFT 학습을 참고하세요.

DROID 작업 적응

실험적: examples/flux3/lora.json은 SO-101 PEFT 옵티마이저, LoRA, dropout, EMA 및 학습 예산 설정을 재사용합니다. 구성 로딩은 확인되었지만, end-to-end DROID PEFT 학습 및 작업 성능은 아직 검증되지 않았습니다.

DROID 정책을 리포지토리 ID로 직접 로드합니다. 공유 인코더는 자동으로 해결됩니다:

lerobot-train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-droid \
  --dataset.repo_id=YOUR_ORG/YOUR_DROID_TASK_DATASET \
  --output_dir=outputs/droid_lora

이 설정은 effective batch 8과 2,500번의 optimizer updates를 사용하며, EMA가 활성화되어 있고 raw/EMA checkpoints가 125번 업데이트마다 저장됩니다. Hub로의 자동 업로드는 비활성화되어 있습니다. 더 큰 예산을 위해 SO-101 예시와 마찬가지로 4개의 GPU로 시작하고 --batch_size=8 --steps=14000을 추가하세요.

프리셋은 DROID 체크포인트의 카메라 레이아웃, 조건부 입력, 액션 표현, 제어율 및 프로세서를 로드하며, 나열된 학습/런타임 설정만 오버라이드합니다. 데이터셋은 체크포인트의 계약과 일치해야 합니다. 현재 패키지는 3개의 카메라, 8개의 상태/액션 채널, 15 Hz의 절대 관절 위치 명령을 사용합니다. 필요 시 --rename_map으로 카메라 이름을 매핑하세요. 두 로봇 모두 동일한 학습 예제를 사용하며, 로드 또는 추론을 위해 로봇별 레시피가 필요하지 않습니다.

기본값 및 정규화

새로운 설정은 이러한 PEFT 기본값을 사용합니다. --policy.path를 로드하면 체크포인트에 저장된 설정이 사용되며, 명시적인 설정/CLI 오버라이드가 우선합니다. DROID 체크포인트는 자체 설정을 유지합니다.

새로운 설정은 n_obs_steps=1, history_snapshots=1, condition_on_past_actions=false를 사용합니다. 체크포인트를 로드하면 명시된 값이 유지됩니다. 고정된 SO-101 패키지는 8개의 히스토리 타임스텝, 2개의 시각 스냅샷, 과거 액션/상태 조건화를 사용합니다. conditioning="history"는 독립적인 스냅샷 인코딩을 선택하며, 타임스텝 수를 단독으로 결정하지는 않습니다. 학습에서는 미래 프레임과 명령을 타깃으로 읽으며, 액션 델타를 계산하기 위해 직전 명령도 읽습니다.

카메라 이름/순서, 상태/행동 차원 및 절대 행동 채널은 로봇별로 다릅니다. SO-101 체크포인트는 마지막 채널을 절대값으로 유지(delta_absolute_dims=[-1])하고 채널 가중치 [1, 1, 1, 1, 1, 2]를 사용합니다. 클래스 기본값은 절대 채널 예외를 사용하지 않고 동일한 가중치를 사용합니다.

베이스의 저장된 프로세서가 정규화 통계를 소유합니다. 원시 상태/행동을 입력하고 학습, 재개, 추론 시 동일한 통계를 유지하세요. SO-101 패키지의 분위수는 해당 패키지의 조인트 순서, 단위, 보정을 전제로 하며, 범용적인 SO-101 보정 값이 아닙니다. 다른 통계를 사용하려면 학습 전에 해당 프로세서와 일치하는 별도의 베이스를 준비하세요. 절대 행동 데이터셋의 통계는 커맨드 델타 통계를 대체할 수 없습니다.

Training은 어댑터와 학습된 헤드를 완전한 base에 대한 참조와 함께 저장합니다. 다시 로드할 때 해당 base와 그 동결된 인코더 파일에 접근할 수 있도록 유지하세요. 어댑터는 base에 자동으로 병합되지 않습니다.

재개

lerobot-train \
  --config_path=outputs/so101_lora/checkpoints/last/pretrained_model/train_config.json \
  --resume=true

pretrained_model 체크포인트에서 재개하여 옵티마이저 및 EMA 상태를 복원합니다. 추론 시 EMA가 활성화된 경우 pretrained_model_ema 형제를 선택할 수 있습니다. 각 어댑터의 원래 베이스와 프로세서를 유지하세요. 새로 다운로드한 베이스는 이전 어댑터에 사용된 것을 대체할 수 없습니다.

추론

LeRobot runner의 --policy.path를 학습된 어댑터 체크포인트로 지정하세요. 이 경우 참조된 베이스와 저장된 프로세서가 로드됩니다. 완전한 베이스를 직접 호출하는 경우, 프로세서를 함께 로드하세요:

from lerobot.policies.factory import make_pre_post_processors
from lerobot.policies.flux3 import Flux3Policy

checkpoint = "black-forest-labs/flux-3-action-so101"
policy = Flux3Policy.from_pretrained(checkpoint).eval()
preprocessor, postprocessor = make_pre_post_processors(policy.config, pretrained_path=checkpoint)

# Run on every control tick, including while actions are queued:
prepared = preprocessor(observation)
normalized_action = policy.select_action(prepared)
absolute_command = postprocessor(normalized_action)

는 관찰/명령 히스토리를 버퍼링합니다. 는 정규화를 되돌리고, 델타를 적분하며, 모든 틱마다 공유 명령 히스토리를 업데이트합니다. 첫 번째 명령은 측정된 상태에 고정되며, 이후의 명령은 이전 명령에서 적분됩니다. 직접적인 정책 출력은 정규화되어 있으며, 반드시 후처리되어야 합니다. 에피소드 경계 또는 개입 이후에는 정책과 두 프로세서를 모두 리셋해야 하며, LeRobot의 동기식 러너는 이를 수행합니다. RTC/비동기 추론은 지원되지 않습니다.

기본 단일 타임스텝 정책의 경우, 현재 이미지와 측정된 상태만 있으면 충분합니다. 명시적인 델타 앵커를 사용한 오프라인 예측의 경우, 이전 명령 하나를 포함하여 observation.command_history를 제공해야 합니다. 이렇게 하면 모델의 조건부에 명령을 추가하지 않고도 출력 적분을 앵커링할 수 있습니다. 여러 관찰 타임스텝으로 명시적으로 구성된 체크포인트에는 policy.config.n_obs_steps개의 이미지/상태와 명령 이력이 필요합니다. 입력을 전처리자를 통해 전달한 후 predict_action_chunk를 호출하고, 후처리를 수행합니다. --policy.compile_model=true는 동결된 인코더와 추론 DiT의 지연 컴파일을 활성화합니다. 학습용 DiT는 컴파일하지 않습니다.

추론 VRAM을 줄이려면, 로딩 후 Qwen을 CPU로 이동하고 마지막 policy.to(...) 호출을 수행하세요:

policy.frozen.text_encoder.cpu()

이 작업은 약 8.3 GiB의 가중치를 호스트 RAM으로 이동합니다. 캐시되지 않은 인스트럭트는 더 느리게 인코딩되지만, 임베딩은 여전히 정책 디바이스로 전송되고 캐시됩니다. 로딩 중 피크 메모리를 줄이지 않으며, 이후의 policy.to(...)는 Qwen을 정책과 함께 다시 이동시킵니다.

다른 구현 방식 학습

아직 파인튜닝된 체크포인트를 제공하지 않는 구현 방식의 경우, flux-action의 베이스 모델에서 시작합니다. 이 워크플로우는 로봇의 액션 헤드와 컨디셔닝 헤드, 카메라 레이아웃, 액션 표현 및 정규화를 정의하는 것을 요구합니다. 차원이 같다는 사실만으로는 다른 로봇의 헤드를 재사용할 수 없습니다.

이후의 LeRobot 추론 또는 작업 적응을 위해, 학습된 모델을 해당 구성 및 일치하는 프로세서와 함께 완전한 LeRobot 정책으로 패키징하세요. 원시 가중치 파일만으로는 로드 가능한 --policy.path 패키지가 아닙니다. 이 가이드는 기존 SO-101 및 DROID 정책을 새로운 작업에 적응시키는 데 초점을 맞춥니다.

모든 정책은 video 및 video_cond 콘텐츠 스트림과 로봇의 액션 스트림만 빌드합니다. 사용되지 않는 이미지/오디오 가중치는 로드 전에 필터링되어, 액션 예측을 변경하지 않으면서 DiT 파라미터의 약 28%를 절약합니다. 저장된 스트림 목록은 이미지/오디오 브랜치를 활성화하지 않으며, 전체 베이스를 로드하고 저장하면 축소된 레이아웃이 기록됩니다. 제거된 스트림을 대상으로 하는 어댑터는 원래 베이스와 이전 코드 리비전이 필요합니다.

완전한 정책은 기본적으로 엄격하게 로드됩니다. 알려진 미사용 스트림을 필터링한 후, 누락된, 예상치 못한, 또는 잘못된 형식의 가중치는 오류를 발생시킵니다. 명시적으로 strict=False를 사용하면 누락된 embodiment head를 초기화하거나 형식이 일치하지 않는 embodiment head의 크기를 조정하는 것만 허용됩니다. trunk 가중치는 여전히 필수이며 형식 검사가 수행됩니다.

검증

CPU 테스트는 packing, 프로세서, PEFT 업데이트 및 체크포인트/재개를 다룹니다. 이전의 실제 가중치 H200 검증은 현재 프로세서 리팩토링보다 먼저 수행되었습니다. 이 정확한 리비전은 여전히 GPU 및 로봇 검증이 필요합니다. 기존 검증은 4-GPU 용량 또는 새로운 embodiment에서의 PEFT 품질을 입증하지 않습니다.