Official code for Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization, published in ACM Transactions on Multimedia Computing, Communications, and Applications (TOMM), 2026.
Jinwen Wang, Youfang Lin, Xiaobo Hu, Qian Xu, Shuo Wang, Zhuo Chen, and Kai Lv
Paper · Environment setup · Implementation and reproduction notes
T2RD separates task-relevant content from task-irrelevant appearance and learns a control policy from the content representation. Its auxiliary learning objective combines representation consistency, cross-reconstruction, and cross-dynamic prediction.
This release contains the DMControl-GB implementation for five tasks. The robotic manipulation experiments in the paper are not included in this repository.
Each auxiliary batch selects Random Convolution or Random Overlay with probability 0.5 each. The current and next observations use the same augmentation type, with independently sampled convolution weights or overlay images. Overlay images come from Places365.
The three auxiliary losses have weights 1, 100, and 0.1, respectively. The critic, actor, and auxiliary representation learner share the content encoder. Style features are used for cross-reconstruction.
The implementation uses the original MuJoCo 2.0 / DMControl-GB environment stack. Please follow setup/README.md for Conda, MuJoCo, EGL rendering, Places365, and evaluation videos. The bundled benchmark environment modifications are required for the color and video tests.
conda create -n t2rd --override-channels -c conda-forge python=3.7.6 pip=23.0.1 setuptools=65.6.3
conda activate t2rd
bash setup/install_envs.sh
bash scripts/download_places365.sh
python scripts/download_benchmark_videos.py
export MUJOCO_GL=eglMuJoCo 2.0 must be installed before install_envs.sh; see the setup guide. Datasets, MuJoCo binaries, license keys, checkpoints, and training logs are not distributed here.
The release follows the authors' actual training command and source defaults. The original --algorithm soda selected the authors' modified T2RD, not the original SODA algorithm. Its unambiguous release equivalent is:
CUDA_VISIBLE_DEVICES=0 python3 src/train.py --algorithm t2rd --seed 1Choose the GPU index available on your machine. --algorithm soda remains a compatibility alias for the same T2RD implementation; it does not select a SODA baseline. The command above uses Finger Spin, 500k training steps, batch size 128, 1,000 initial random steps, auxiliary learning rate 1e-3, and the author-confirmed 50/50 convolution/overlay augmentation with all three losses.
Evaluation runs automatically during training every 20k steps (10k for Ball-in-Cup), at episode boundaries. By default, the first evaluation is at the first interval, matching the local training script's step > 0 condition. Add --eval_at_start to include an untrained step-0 baseline, as present in some older runs.
Three separate, persistent test environments are created for color_hard, video_easy, and video_hard. Every evaluation round evaluates all three sequentially with 30 episodes each; they are not alternated across rounds. The original training-domain evaluation is also retained. Download the benchmark videos before starting. Use --eval_freq and --eval_episodes to select the intended historical schedule; both 10k and 20k intervals occur in archived runs. --eval_freq 0 disables automatic evaluation for an installation-only smoke test.
Use --task to select the task and its action repeat together. This replaces manually editing domain_name, task_name, and action_repeat in the old arguments.py:
| Task | --task |
Action repeat |
|---|---|---|
| Finger Spin | finger_spin |
2 |
| Cartpole Swingup | cartpole_swingup |
8 |
| Walker Walk | walker_walk |
4 |
| Walker Stand | walker_stand |
4 |
| Ball in Cup Catch | ball_in_cup_catch |
4 |
For the historical Ball-in-Cup configuration:
CUDA_VISIBLE_DEVICES=0 python3 src/train.py --algorithm t2rd \
--task ball_in_cup_catch --seed 1 --eval_freq 10k --eval_at_startFor another task:
CUDA_VISIBLE_DEVICES=0 python src/train.py \
--task walker_walk --seed 0 --train_steps 500k \
--augmentation mixed --overlay_probability 0.5 --aux_lr 1e-3scripts/train.sh uses the same defaults. The source command is authoritative for this release; manuscript settings are not substituted for the values used by that command. Use a different --log_dir for different configurations.
Runs are saved under logs/<task>/<algorithm>/<seed>/ (t2rd by default). They include configuration (info.json), training metrics (train.log), automatic evaluation metrics (eval.log), per-evaluation episode rewards (evaluation/*.json), periodic checkpoints, and the final training state (model/final.pt). Periodic saves retain the source's episode-boundary schedule. --save_video records the first episode of each evaluation setting.
For a short installation check after downloading Places365:
CUDA_VISIBLE_DEVICES=0 python src/train.py --task finger_spin --seed 0 \
--train_steps 12 --init_steps 4 --batch_size 4 --eval_freq 0 --log_dir logs/smokeTraining already evaluates automatically. To evaluate an existing checkpoint again, run the 500k checkpoint over 30 episodes in all three visual generalization settings:
CUDA_VISIBLE_DEVICES=0 python src/eval.py --seed 0 \
--checkpoint logs/finger_spin/t2rd/0/model/500000.pt \
--eval_modes color_hard video_easy video_hard --eval_episodes 30Use --eval_modes train for the training visual domain and --save_video to record the first episode. Evaluation reads the task and architecture from the checkpoint and saves per-episode rewards, their mean, and their standard deviation as JSON. This within-run standard deviation is not the paper's standard deviation across training seeds.
For the smoke checkpoint, pass logs/smoke/finger_spin/t2rd/0/model/final.pt and --eval_episodes 1. A short smoke run checks execution, not learned performance.
src/algorithms/t2rd.py Three auxiliary objectives and their update schedule
src/algorithms/networks.py Content/style encoder, policy, critic, reconstruction, dynamics
src/algorithms/sac.py Soft Actor-Critic updates
src/augmentations.py Random crop, convolution, and Places365 overlay
src/train.py Environment interaction and training
src/eval.py Optional evaluation of an existing checkpoint
src/evaluation.py Shared evaluation used during training and offline
src/checkpoint.py State-dict serialization
src/env/ Original modified DMControl-GB environment dependencies
setup/ Environment setup and dataset configuration
scripts/ Training, evaluation, and data download commands
tests/ Numerical comparison and smoke checks
The release removes unused experimental losses while retaining the active computation and necessary BatchNorm state updates. The inactive 50/50 augmentation selector in the supplied active entry point has been restored. Implementation notes describe the original files, retained numerical conventions, and the exact scope of validation.
Numerical agreement under matched weights and random state does not establish reproduction of the paper's full learning curves. Removing unused network constructors changes random-number consumption during initialization; the same seed alone is not a guarantee of identical old and new runs. The repository does not distribute pretrained weights or newly reproduced benchmark scores.
The training infrastructure and benchmark derive from DMControl Generalization Benchmark / SODA, with contributions from SAC-AE, dmc2gym, and DeepMind Control Suite. We retain their license notices; see LICENSE and THIRD_PARTY_NOTICES.md. External datasets retain their own terms.
If you use this code or find our work helpful, please cite:
@article{wang2026t2rd,
author = {Wang, Jinwen and Lin, Youfang and Hu, Xiaobo and Xu, Qian
and Wang, Shuo and Chen, Zhuo and Lv, Kai},
title = {Task-Relevant Representation Decoupling for Visual
Reinforcement Learning Generalization},
journal = {ACM Transactions on Multimedia Computing, Communications,
and Applications},
year = {2026},
volume = {22},
number = {8},
articleno = {211},
numpages = {24},
month = jul,
doi = {10.1145/3820367},
url = {https://doi.org/10.1145/3820367}
}