{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/improving-transformer-world-models-for-data","title":"Improving Transformer World Models for Data-Efficient RL","arxiv_id":"2502.01591","date":"2025-02-03","proceeding":null,"authors":["Antoine Dedieu","Joseph Ortiz","Xinghua Lou","Carter Wendelken","Wolfgang Lehrach","J Swaroop Guntupalli","Miguel Lazaro-Gredilla","Kevin Patrick Murphy"],"abstract":"We present an approach to model-based RL that achieves a new state of the art performance on the challenging Craftax-classic benchmark, an open-world 2D survival game that requires agents to exhibit a wide range of general abilities -- such as strong generalization, deep exploration, and long-term reasoning. With a series of careful design choices aimed at improving sample efficiency, our MBRL algorithm achieves a reward of 67.4% after only 1M environment steps, significantly outperforming DreamerV3, which achieves 53.2%, and, for the first time, exceeds human performance of 65.0%. Our method starts by constructing a SOTA model-free baseline, using a novel policy architecture that combines CNNs and RNNs. We then add three improvements to the standard MBRL setup: (a) \"Dyna with warmup\", which trains the policy on real and imaginary data, (b) \"nearest neighbor tokenizer\" on image patches, which improves the scheme to create the transformer world model (TWM) inputs, and (c) \"block teacher forcing\", which allows the TWM to reason jointly about the future tokens of the next timestep.","url_abs":"https://arxiv.org/abs/2502.01591v1","url_pdf":"https://arxiv.org/pdf/2502.01591v1.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[],"tasks":[],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2502.01591","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2502.01591"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"deterministic:regex_extraction","url":"https://github.com/eloialonso/iris","reach":null},{"provenance":"deterministic:regex_extraction","url":"https://github.com/luchris429/purejaxrl","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"ran":2,"ran_draft_wrong":6,"ran_fixture":1,"unverified":4},"by_repo_kind":{"found_in_text":{"samples":13,"ran":9,"repositories":2}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":10,"samples":[{"code_sha256_prefix":"878ced2dad28f2bb","entry":"ScalingLayer","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":"invariant","behaviour_fingerprint":true,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"878ced2dad28f2bb"}},{"code_sha256_prefix":"970bb08b535c4c9e","entry":"get_ckpt_path","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"970bb08b535c4c9e"}},{"code_sha256_prefix":"ec8a321400aae8d2","entry":"init_VinvB","repo":"luchris429/purejaxrl","repo_kind":"found_in_text","path":"purejaxrl/experimental/s5/s5.py","file_url":"https://github.com/luchris429/purejaxrl/blob/HEAD/purejaxrl/experimental/s5/s5.py","link_basis":"harvester_set","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":"invariant","behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ec8a321400aae8d2"}},{"code_sha256_prefix":"a4bdc24302b8ab98","entry":"init_log_steps","repo":"luchris429/purejaxrl","repo_kind":"found_in_text","path":"purejaxrl/experimental/s5/s5.py","file_url":"https://github.com/luchris429/purejaxrl/blob/HEAD/purejaxrl/experimental/s5/s5.py","link_basis":"harvester_set","language":"python","status":"ran_fixture","verification_level":1,"contract_check":"RAISES","metamorphic_tier":"invariant","behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a4bdc24302b8ab98"}},{"code_sha256_prefix":"8b1a3bc2f6b64dce","entry":"log_step_initializer","repo":"luchris429/purejaxrl","repo_kind":"found_in_text","path":"purejaxrl/experimental/s5/s5.py","file_url":"https://github.com/luchris429/purejaxrl/blob/HEAD/purejaxrl/experimental/s5/s5.py","link_basis":"harvester_set","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"8b1a3bc2f6b64dce"}},{"code_sha256_prefix":"ebaf93caaa927ff9","entry":"md5_hash","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"ebaf93caaa927ff9"}},{"code_sha256_prefix":"4d9196e5759e6950","entry":"normalize_tensor","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"4d9196e5759e6950"}},{"code_sha256_prefix":"6855e475947477db","entry":"spatial_average","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"6855e475947477db"}},{"code_sha256_prefix":"ca3575bced071dc5","entry":"vgg16","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"ca3575bced071dc5"}},{"code_sha256_prefix":"4767c4ed7a66b542","entry":"LPIPS","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"4767c4ed7a66b542"}},{"code_sha256_prefix":"8681e7e0f15169da","entry":"NetLinLayer","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"8681e7e0f15169da"}},{"code_sha256_prefix":"a823cb01993404cf","entry":"Tokenizer","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"a823cb01993404cf"}},{"code_sha256_prefix":"5b78fedd249dc684","entry":"download","repo":"eloialonso/iris","repo_kind":"found_in_text","path":"src/models/tokenizer/tokenizer.py","file_url":"https://github.com/eloialonso/iris/blob/HEAD/src/models/tokenizer/tokenizer.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"GPL-3.0","inline_ok":false,"mcp_get_code":{"code_sha256":"5b78fedd249dc684"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}