DexForce · yuecideng · Mar 7, 2026 · Feb 27, 2026 · Feb 28, 2026 · Feb 28, 2026
diff --git a/CLAUDE.md b/CLAUDE.md
@@ -80,7 +80,7 @@ Every source file begins with the Apache 2.0 copyright header:
 - Use full type hints on all public APIs.
 - Use `from __future__ import annotations` at the top of every file.
 - Use `TYPE_CHECKING` guards for circular-import-safe imports.
-- Prefer `Union[A, B]` or `A | B` (Python 3.10+ union syntax is acceptable).
+- Prefer `A | B` over `Union[A, B]`.
 
 ### Configuration Pattern (`@configclass`)
 
@@ -265,7 +265,7 @@ def test_edge_case():
     assert result is not None
 ```
 
-**`unittest.TestCase` style** — when tests must run in a specific order or share `setUp`/`tearDown` state:
+**`Class` style** — when tests must run in a specific order or share `setup_method`/`teardown_method` state:
 
 ```python
 # ----------------------------------------------------------------------------
@@ -274,28 +274,23 @@ def test_edge_case():
 # ...
 # ----------------------------------------------------------------------------
 
-import unittest
 from embodichain.my_module import MyClass
 
 
-class TestMyClass(unittest.TestCase):
-    def setUp(self):
+class TestMyClass():
+    def setup_method(self):
         self.obj = MyClass(param=1.0)
 
-    def tearDown(self):
+    def teardown_method(self):
         pass
 
     def test_basic_behavior(self):
         result = self.obj.run()
-        self.assertEqual(result, expected)
+        assert result == expected_result
 
     def test_raises_on_bad_input(self):
-        self.assertRaises(ValueError, self.obj.run, bad_input)
-
-
-if __name__ == "__main__":
-    unittest.main()
-```
+        with pytest.raises(ValueError):
+            self.obj.run(bad_input)
 
 ### Conventions
 

diff --git a/configs/agents/rl/basic/cart_pole/gym_config.json b/configs/agents/rl/basic/cart_pole/gym_config.json
@@ -1,6 +1,7 @@
 { 
     "id": "CartPoleRL",
     "max_episodes": 5,
+    "max_episode_steps": 500,
     "env": {
         "events": {},
         "observations": {
@@ -26,7 +27,6 @@
         },
         "extensions": {
             "action_type": "delta_qpos",
-            "episode_length": 500,
             "action_scale": 0.1,
             "success_threshold": 0.1
         }

diff --git a/configs/agents/rl/push_cube/gym_config.json b/configs/agents/rl/push_cube/gym_config.json
@@ -1,6 +1,7 @@
 { 
     "id": "PushCubeRL",
     "max_episodes": 5,
+    "max_episode_steps": 100,
     "env": {
         "events": {
             "randomize_cube": {
@@ -112,7 +113,6 @@
         },
         "extensions": {
             "action_type": "delta_qpos",
-            "episode_length": 100,
             "action_scale": 0.1,
             "success_threshold": 0.1
         }

diff --git a/configs/gym/agent/pour_water_agent/fast_gym_config.json b/configs/gym/agent/pour_water_agent/fast_gym_config.json
@@ -251,15 +251,15 @@
                 "mode": "save",
                 "params": {
                     "robot_meta": {
-                        "control_freq": 25,
-                        "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
+                        "control_freq": 25
                     },
                     "instruction": {
                         "lang": "Pour water from the bottle into the mug."
                     }
                 }
             }
         },
+        "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
         "success_params": {
             "strict": false
         }

diff --git a/configs/gym/agent/rearrangement_agent/fast_gym_config.json b/configs/gym/agent/rearrangement_agent/fast_gym_config.json
@@ -236,15 +236,15 @@
                 "mode": "save",
                 "params": {
                     "robot_meta": {
-                        "control_freq": 25,
-                        "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
+                        "control_freq": 25
                     },
                     "instruction": {
                         "lang": "Place the spoon and fork neatly into the plate on the table."
                     }
                 }
             }
         },
+        "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
         "success_params": {
             "strict": false
         }

diff --git a/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json b/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json
@@ -104,7 +104,7 @@
                 "mode": "modify",
                 "name": "robot/qpos",
                 "params": {
-                    "joint_ids": [12, 13, 14, 15]
+                    "joint_ids": [6, 13]
                 }
             },
             "block_1_pose": {

diff --git a/configs/gym/blocks_ranking_size/cobot_magic_3cam.json b/configs/gym/blocks_ranking_size/cobot_magic_3cam.json
@@ -78,7 +78,7 @@
                 "mode": "modify",
                 "name": "robot/qpos",
                 "params": {
-                    "joint_ids": [12, 13, 14, 15]
+                    "joint_ids": [6, 13]
                 }
             },
             "block_1_pose": {

diff --git a/configs/gym/match_object_container/cobot_magic_3cam.json b/configs/gym/match_object_container/cobot_magic_3cam.json
@@ -61,7 +61,7 @@
                 "mode": "modify",
                 "name": "robot/qpos",
                 "params": {
-                    "joint_ids": [12, 13, 14, 15]
+                    "joint_ids": [6, 13]
                 }
             },
             "block_cube_1_pose": {

diff --git a/configs/gym/pour_water/gym_config.json b/configs/gym/pour_water/gym_config.json
@@ -1,6 +1,7 @@
 {   
     "id": "PourWater-v3",
     "max_episodes": 10,
+    "max_episode_steps": 300,
     "env": {
         "events": {
             "random_light": {
@@ -200,7 +201,7 @@
                 "mode": "modify",
                 "name": "robot/qpos",
                 "params": {
-                    "joint_ids": [12, 13, 14, 15]
+                    "joint_ids": [6, 13]
                 }
             },
             "bottle_pose": {
@@ -264,8 +265,7 @@
                 "params": {
                     "robot_meta": {
                         "robot_type": "CobotMagic",
-                        "control_freq": 25,
-                        "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
+                        "control_freq": 25
                     },
                     "instruction": {
                         "lang": "Pour water from bottle to cup"
@@ -278,7 +278,8 @@
                     "use_videos": true
                 }
             }
-        }
+        },
+        "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
     },
     "robot": {
         "uid": "CobotMagic",

diff --git a/configs/gym/pour_water/gym_config_simple.json b/configs/gym/pour_water/gym_config_simple.json
@@ -1,6 +1,7 @@
 {   
     "id": "PourWater-v3",
     "max_episodes": 5,
+    "max_episode_steps": 300,
     "env": {
         "events": {
             "record_camera": {
@@ -213,8 +214,7 @@
                 "params": {
                     "robot_meta": {
                         "robot_type": "CobotMagic",
-                        "control_freq": 25,
-                        "control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
+                        "control_freq": 25
                     },
                     "instruction": {
                         "lang": "Pour water from bottle to cup"

diff --git a/configs/gym/stack_blocks_two/cobot_magic_3cam.json b/configs/gym/stack_blocks_two/cobot_magic_3cam.json
@@ -41,7 +41,7 @@
                 "mode": "modify",
                 "name": "robot/qpos",
                 "params": {
-                    "joint_ids": [12, 13, 14, 15]
+                    "joint_ids": [6, 13]
                 }
             },
             "block_1_pose": {

diff --git a/configs/gym/stack_cups/cobot_magic_3cam.json b/configs/gym/stack_cups/cobot_magic_3cam.json
@@ -41,7 +41,7 @@
                 "mode": "modify",
                 "name": "robot/qpos",
                 "params": {
-                    "joint_ids": [12, 13, 14, 15]
+                    "joint_ids": [6, 13]
                 }
             },
             "cup_1_pose": {

diff --git a/docs/source/api_reference/embodichain/embodichain.agents.rst b/docs/source/api_reference/embodichain/embodichain.agents.rst
@@ -7,6 +7,48 @@
 
    .. autosummary::
 
-      dexforce_vla
+      datasets
+      engine
       rl
 
+Datasets
+--------
+
+.. automodule:: embodichain.agents.datasets
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+   .. autosummary::
+
+      online_data
+      sampler
+
+Online Data Engine
+------------------
+
+.. automodule:: embodichain.agents.engine
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+   .. autosummary::
+
+      data
+
+Reinforcement Learning
+----------------------
+
+.. automodule:: embodichain.agents.rl
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+   .. autosummary::
+
+      algo
+      buffer
+      models
+      train
+      utils
+
diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.rst
@@ -89,22 +89,20 @@ Registration System
 
     :param name: Unique identifier for the environment
     :param cls: Environment class (must inherit from BaseEnv or BaseEnv)  
-    :param max_episode_steps: Maximum steps per episode (optional)
     :param default_kwargs: Default keyword arguments for environment creation
 
 .. autofunction:: register_env
 
     Decorator function for registering environment classes. This is the recommended way to register environments.
 
     :param uid: Unique identifier for the environment
-    :param max_episode_steps: Maximum steps per episode (optional)
     :param override: Whether to override existing environment with same ID
     :param kwargs: Additional registration parameters
 
     Example:
         .. code-block:: python
 
-            @register_env("MyEnv-v1", max_episode_steps=1000)
+            @register_env("MyEnv-v1")
             class MyCustomEnv(BaseEnv):
                 def __init__(self, **kwargs):
                     super().__init__(**kwargs)

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -29,6 +29,7 @@ Table of Contents
 
    overview/sim/index
    overview/gym/index
+   overview/agents/online_data.md
    overview/rl/index
 
 .. toctree::