bj.choi - notes

[논문 리뷰] TMax로 이해하는 Terminal-Agent RL과 DPPO

Updated:

TMax를 통해 terminal-agent RL에서 task prompt, sandbox, verifier, GRPO/DPPO가 어떻게 연결되는지 정리합니다.

1. Why TMax

처음 고민은 TMax가 아니라 Codex 채팅 데이터였습니다.

Codex에는 질의응답만 남는 것이 아니라, 사용자 요청, tool call, terminal output, 실패 해석, 다음 행동, 검증 결과가 이어지는 작업 흐름이 남았는데요. 이 데이터를 그대로 백업하는 것과 별개로, 나중에 모델 학습에 활용하려면 어떤 구조로 봐야 하는지와 어떻게 학습할 수 있는지가 궁금했습니다.

처음에는 Codex 로그를 SFT 데이터처럼 사용자 요청과 이상적인 응답으로 바꾸는 방향을 생각했습니다. 하지만 Codex 데이터는 최종 답변만 잘하는 모델이 아니라, 실제 환경에서 명령을 실행하고 결과를 읽고 다음 행동을 정해야 하는 모델입니다.

이 때문에 다른 모델이나 유사한 작업을 수행하는 모델을 만들 때 데이터를 어떻게 만들고 학습하는지 확인할 필요가 있었습니다.

그 과정에서 TMax를 살펴봤습니다.

TMax는 공개 모델만 있는 것이 아니라, 논문, 학습 데이터셋, GitHub 코드가 함께 정리되어 있습니다. Hugging Face 데이터셋에는 messages, ground_truth, env_config가 있고, GitHub 코드에는 sandbox, verifier, RL 학습 루프가 연결되어 있습니다.

이 글은 TMax를 성능 벤치마크로 평가하려는 글이 아닙니다.

Codex 작업 로그를 학습 데이터 관점에서 어떻게 바라볼지 고민하면서, TMax가 terminal-agent 학습을 어떤 구조로 풀었는지 확인한 정리입니다. 특히 task prompt + sandbox + verifier + GRPO/DPPO 구조가 무엇을 의미하는지 이해하는 것이 목적입니다.

2. TMax 논문 요약

논문 기준으로 TMax는 terminal agent를 만들기 위한 단순한 레시피를 제안합니다. 구성은 크게 두 가지입니다.

제가 이해한 정의는 이렇습니다.

  1. TMax는 정답 응답을 맞추는 instruction fine-tuning이 아닙니다.
  2. terminal task를 여러 번 풀게 하고 verifier reward를 기준으로 정책을 업데이트하는 terminal-agent RL 학습 방식입니다.

여기서 핵심은 “터미널 명령어를 생성한다”가 아닙니다.
모델이 만든 명령들이 실제 환경 상태를 바꾸고, 그 최종 상태가 task 요구사항을 만족하는지 확인한다는 점입니다.

터미널 모델 관점에서 TMax의 학습 대상은 이렇게 풀어볼 수 있습니다.

3. TMax 데이터셋에는 무엇이 들어 있나

공개 데이터셋인 allenai/tmax-15k-open-instruct는 약 14.6k개의 task row를 갖습니다. 주요 컬럼은 다음과 같습니다.

messages
ground_truth
dataset
env_config
source

여기서 messages는 모델에게 주어지는 task prompt입니다. ground_truth는 보통 task id 역할을 합니다. env_config에는 어떤 sandbox 환경을 띄울지, 어떤 image와 task id를 쓸지가 들어 있습니다.

핵심은 이 데이터셋이 완성된 multi-turn trajectory 모음이 아니라는 데 있습니다. 데이터셋은 시작 조건을 줍니다. 실제 trajectory는 학습 중 현재 policy가 새로 만듭니다.

흐름은 이렇게 볼 수 있습니다.

HF dataset row
-> task prompt와 env_config를 읽음
-> sandbox reset
-> 모델이 bash tool을 사용하며 trajectory 생성
-> submit
-> verifier 실행
-> reward 생성

TMax 데이터셋을 볼 때 messages만 있다고 해서 “멀티턴 데이터가 없다”고 보면 안 됩니다. 멀티턴은 저장된 것이 아니라 rollout 과정에서 생성됩니다.

4. Trajectory, Rollout, Sandbox, Verifier

TMax를 이해하려면 용어를 먼저 분리해야 합니다.

trajectory
= task 시작부터 submit/reward까지 이어지는 전체 풀이 경로

rollout
= 현재 policy가 한 번 샘플링해서 만든 trajectory

agent harness
= 모델 출력과 bash tool 실행 사이를 중개하는 adapter 계층

sandbox
= 모델이 낸 명령이 실제 실행되는 격리 환경

verifier
= submit 후 최종 상태를 검사해 reward를 만드는 채점기

  1. trajectory는 명령 실행, 출력 관찰, 파일 수정, 테스트 실행, submit까지 포함한 전체 상호작용 기록입니다.

  2. agent harness는 모델이 낸 출력을 tool call로 해석하고, 실행 결과를 다시 observation으로 돌려줍니다.

  1. sandbox는 실행 환경입니다.

  2. 검증은 verifier가 합니다.

모델이 명령을 실행하면 그 결과가 sandbox 안의 파일, 프로세스, 출력 상태로 남습니다.
submit 이후 verifier가 그 최종 상태를 확인해서 평가합니다.

5. Verifier 평가 방식

TMax에서 verifier는 모델이 “해결했습니다”라고 말했는지를 보지 않습니다. 최종 환경 상태가 task 요구사항과 맞는지를 봅니다.

예를 들어 task가 버그 수정이라면 verifier는 이런 것을 확인할 수 있습니다.

TMax식 terminal-agent RL에서는 verifier 설계에서 성패가 갈립니다. verifier가 너무 느슨하면 모델은 task를 푸는 대신 verifier를 통과하는 편법을 배울 수 있습니다. 반대로 너무 좁으면 올바른 다른 풀이까지 실패로 처리할 수 있습니다.

제가 이해한 좋은 verifier의 최소 조건은 다음과 같습니다.

TMax의 verifier kind는 taxonomy 기준으로 다섯 가지입니다.

exact_text
metric_threshold
adversarial_corpus
fuzz_equivalence
multi_protocol

다만 이것은 고정 함수 다섯 개가 아닙니다. 실제 verifier는 task마다 만들어진 테스트 코드이고, 위 다섯 가지는 verifier를 어떤 방식으로 만들지 정하는 유형에 가깝습니다.

6. Why Prompt trajectory?

GRPO와 DPPO를 이해하려면 같은 task prompt에 대해 rollout을 여러 개 만드는 이유부터 봐야 합니다.

TMax 설정에서는 보통 한 prompt에 대해 여러 trajectory를 생성합니다. 예를 들어 group size가 32라면 구조는 이렇습니다.

task prompt 1개
-> rollout/trajectory 32개
-> reward 32개
-> advantage 32개

reward는 trajectory마다 하나씩 붙습니다.
[1, 0, 1, 0, ...] 같은 값은 reward 하나가 아니라, 같은 prompt에서 나온 reward 리스트입니다.

이렇게 여러 rollout을 만드는 이유는 비교하기 위해서입니다.

[0, 0, 0, 0] -> 모두 실패. 무엇을 강화해야 할지 알기 어렵다.
[1, 1, 1, 1] -> 모두 성공. 이미 너무 쉬운 문제일 수 있다.
[0, 1, 0, 1] -> 성공과 실패가 섞여 있어 비교할 수 있다.

좋은 RL task는 같은 prompt에서 여러 rollout을 만들었을 때 reward variance가 있어야 합니다.
일부 trajectory는 성공하고 일부 trajectory는 실패해야 무엇을 강화하고 무엇을 약화할지 배울 수 있습니다.

이 관점에서 TMax의 데이터 설계는 “학습 신호가 생기는 난이도와 verifier를 섞는다”에 가깝습니다.
좋은 데이터를 만들기 위한 과정입니다.

7. Reward와 Advantage를 구분하기

처음에는 reward를 계산하고 그대로 학습하면 될 것처럼 보였습니다.
하지만 GRPO에서는 reward를 그대로 쓰기보다 advantage를 계산해야 합니다.

제가 이해한 구분을 정리했습니다.

reward
= trajectory 하나에 붙는 절대 점수

advantage
= 같은 prompt group 안에서 내 reward가 평균보다 얼마나 좋은지/나쁜지

예를 들어 두 prompt group이 있습니다.

Prompt A rewards = [1, 1, 1, 0]
mean A = 0.75

reward 1 advantage = 1 - 0.75 = +0.25
reward 0 advantage = 0 - 0.75 = -0.75
Prompt B rewards = [0, 0, 0, 1]
mean B = 0.25

reward 1 advantage = 1 - 0.25 = +0.75
reward 0 advantage = 0 - 0.25 = -0.25

둘 다 reward 1이지만 의미가 다릅니다.
쉬운 task에서 성공한 것은 평균보다 조금 나은 것이고, 어려운 task에서 성공한 것은 평균보다 많이 나은 것입니다.

반대로 reward 0도 의미가 다릅니다. 쉬운 task에서 실패하면 크게 약화됩니다. 어려운 task에서 실패하면 덜 약화됩니다.

여기서 봐야 할 부분은 평균을 batch 전체에서 내지 않는다는 점입니다. 같은 prompt task에서 나온 여러 trajectory의 reward 평균으로 advantage를 계산하고, 그 advantage 방향으로 모델 정책을 업데이트합니다.

Prompt A의 32개 reward 안에서 평균 계산
Prompt B의 32개 reward 안에서 평균 계산
...
Prompt H의 32개 reward 안에서 평균 계산

서로 다른 prompt의 reward를 직접 비교하면 안 됩니다. task 난이도가 다르기 때문입니다.

8. GRPO는 Critic 없이 Group 평균으로 Advantage를 만든다

PPO 계열 RL에서는 보통 value model이나 critic이 필요합니다. critic은 현재 상태나 답변이 어느 정도 reward를 받을 만한지 예측합니다. 실제 reward와 예측값의 차이를 이용해 advantage를 계산합니다.

GRPO는 여기서 다르게 갑니다.

별도 value model을 두지 않고, 같은 prompt에서 여러 rollout을 만든 뒤 그 group의 reward 평균을 기준선처럼 씁니다.

제가 이해한 GRPO는 이렇습니다.

GRPO는 같은 prompt에 대한 여러 trajectory를 verifier로 채점하고, group reward 평균보다 잘한 trajectory는 강화하며 못한 trajectory는 약화하는 방식입니다.

여기서 학습되는 것은 reward 숫자 자체가 아닙니다.
reward를 받은 trajectory 안의 모델 출력 토큰 확률이 advantage 방향으로 업데이트됩니다.

  1. 예를 들어 reward 1 trajectory가 관련 파일을 탐색하고 올바른 코드를 수정하고 테스트를 실행했다고 가정합니다.
  2. 그 trajectory를 구성한 토큰들의 확률이 올라갑니다.
  3. reward 0 trajectory가 엉뚱한 파일을 수정하거나 테스트 없이 submit했다면, 그 trajectory의 토큰 확률은 내려갑니다.

물론 모델은 어느 한 명령이 정확히 실패 원인인지 직접 알지는 못합니다. sparse reward 학습은 어렵습니다.
GRPO는 여러 rollout을 비교하면서 성공 trajectory 전체를 강화하고 실패 trajectory 전체를 약화하는 쪽으로 접근합니다.

9. TMax에서 DPPO 구성

GRPO까지만 보면 “여러 rollout의 reward를 상대 비교해서 업데이트한다”입니다.
그런데 TMax는 GRPO만 쓰지 않고 DPPO를 사용합니다.

이유는 terminal-agent trajectory가 길고, rollout과 training 사이의 logprob mismatch가 커질 수 있기 때문입니다.

TMax 학습에서는 rollout을 vLLM 같은 inference 엔진에서 생성합니다. 이후 trainer가 같은 trajectory에 대해 training model 기준 logprob를 다시 계산합니다. 이때 둘이 너무 다르면 문제가 됩니다.

  1. 예를 들어 rollout 당시에는 어떤 토큰의 확률이 높았는데, training 때 다시 계산하니 매우 낮게 나온다고 합시다.
  2. 그러면 그 토큰을 현재 모델의 안정적인 행동처럼 취급해 강하게 업데이트하는 것이 위험할 수 있습니다.
  3. DPPO는 여기서 mask를 사용합니다.
rollout logprob와 training logprob가 너무 다른 토큰
-> 이번 loss/update 계산에서 제외

핵심은 이것이 token ban이 아니라 loss mask라는 점입니다. 토큰 자체가 나쁜 것이 아닐 수 있습니다.
문제는 이번 trajectory에서 rollout 시점과 training 시점의 확률 계산이 너무 달라, 그 토큰을 학습 신호로 믿기 어렵다는 점입니다.

DPPO는 GRPO 위에 안정화 장치를 추가한 것으로 볼 수 있습니다.

DPPO는 GRPO처럼 group reward로 advantage를 계산하되, rollout logprob와 training logprob가 크게 어긋나는 토큰을 mask해서 긴 terminal-agent trajectory에서 업데이트 불안정을 줄이는 방법입니다.

10. TMax 학습 루프를 한 번에 보면

지금까지의 개념을 하나의 루프로 보면 다음과 같습니다.

1. task prompt를 batch로 고른다.
2. 각 prompt마다 여러 rollout을 만든다.
3. 각 rollout은 하나의 trajectory가 된다.
4. agent harness가 모델 출력을 bash tool call로 파싱한다.
5. sandbox에서 bash 명령이 실행되고 observation이 모델에게 돌아간다.
6. submit하면 verifier가 최종 상태를 검사해 reward를 만든다.
7. 같은 prompt의 reward group으로 advantage를 계산한다.
8. DPPO loss가 logprob mismatch가 큰 토큰을 mask한다.
9. trainer가 policy weight를 업데이트한다.

TMax 설정에서 자주 나오는 숫자는 다음처럼 이해하면 됩니다.

unique prompts per rollout batch = 8
group size = 32

그러면 한 batch에는 다음이 생깁니다.

8 prompts * 32 rollouts = 256 trajectories
256 rewards
256 advantages

하지만 advantage는 전체 256개의 reward 평균으로 계산하지 않습니다. 각 prompt별 32개 reward 안에서 계산합니다.

11. 내가 직접 실습한다면 무엇부터 볼까

TMax 전체를 바로 재현하려고 하면 범위가 큽니다. 모델, vLLM rollout, sandbox pool, verifier, Docker/Podman, DPPO loss, checkpoint, evaluation이 한꺼번에 얽힙니다.

첫 실습 목표는 TMax-4B 성능 재현이 아니라, DPPO 학습 루프가 실제로 한 번 도는 것을 확인하는 쪽이 더 현실적입니다.

제가 잡은 첫 목표는 이렇습니다.

RunPod 1GPU 환경에서 작은 Qwen 계열 모델과 TMax debug script를 사용해 rollout -> verifier reward -> group advantage -> DPPO mask -> trainer update 흐름이 실제로 동작하는지 점검합니다.

실습 전에 코드에서 먼저 볼 파일은 다음입니다.

training/open-instruct/open_instruct/grpo_fast.py
training/open-instruct/open_instruct/environments/swerl_vanillux_sandbox.py
training/open-instruct/scripts/tmax/RL/qwen35_2b_1gpu.sh

확인할 질문은 단순합니다.

이 질문에 답할 수 있으면, 단순히 스크립트를 실행한 것이 아니라 TMax 학습 루프를 이해하면서 실행한 것이 됩니다.

12. 결론

TMax는 terminal model을 모델이 실제 환경에서 행동하고 그 행동이 만든 최종 상태를 verifier가 채점하며, 그 reward를 바탕으로 policy를 업데이트하는 문제로 봅니다.

GRPO는 같은 prompt에서 여러 trajectory를 만들고 group reward 평균을 기준으로 advantage를 계산합니다.
별도 value model 없이도 상대적으로 좋은 trajectory와 나쁜 trajectory를 구분할 수 있습니다.

DPPO는 여기에 안정화 장치를 더합니다. rollout 때와 training 때 logprob가 크게 다른 토큰은 이번 update에서 mask합니다.
긴 terminal trajectory에서는 이런 mismatch가 학습 붕괴로 이어질 수 있기 때문에, TMax는 DPPO와 FP32 LM head, 큰 group size 같은 설정을 함께 사용합니다.

제가 이해한 TMax의 핵심은 다음 문장으로 정리할 수 있습니다.

TMax는 정답 trajectory를 외우게 하는 모델이 아니라, terminal task를 직접 여러 번 풀고 verifier reward로 자신의 행동 확률을 조정하는 모델을 만드는 레시피입니다.

13. Codex에서 활용 방식

Codex를 오래 사용하면 단순 채팅 로그가 아니라 작업 단위의 기록이 쌓였습니다. 사용자 요청, 중간 명령, 실패한 실행 결과, 수정 방향, 최종 검증까지 남기 때문에 TMax식 학습 데이터 설계에 필요한 단서를 얻을 수 있습니다.

이 로그를 바로 학습 데이터로 쓰기보다는 먼저 task prompt와 verifier 후보를 분리하는 작업이 필요합니다.
어떤 요청이 하나의 prompt task가 될 수 있는지, 어떤 파일 변경이나 테스트 결과를 최종 상태로 검증할 수 있는지부터 정리해야 합니다.

소형 모델도 이미 안정적인 터미널 작업 일부는 수행할 수 있는 수준에 와 있습니다. Codex 로그를 바탕으로 로컬에서 활용할 수 있는 작은 terminal agent를 구성하는 실험은 앞으로 더 현실적인 방향이 될 것 같습니다.

References