diff --git a/docs/contrib.md b/docs/contrib.md
index 70b087735..6a150d696 100644
--- a/docs/contrib.md
+++ b/docs/contrib.md
@@ -1,15 +1,18 @@
 # Contribution Guide
 
-Thank you for your interest in contributing to AReaL! We welcome contributions from everyone, whether you're fixing bugs, improving documentation, or adding new system and algorithmic features.
+Thank you for your interest in contributing to AReaL! We welcome contributions from
+everyone, whether you're fixing bugs, improving documentation, or adding new system and
+algorithmic features.
 
 ## Setting Up Your Development Environment
 
-New contributors do not have write permissions to the official repository. Please fork the repository and clone your fork locally. AReaL is fully Python-based, making installation straightforward.
+New contributors do not have write permissions to the official repository. Please fork
+the repository and clone your fork locally. AReaL is fully Python-based, making
+installation straightforward.
 
 ```bash
 git clone https://github.com/${your-username}/AReaL
 cd AReaL
-pip3 install -r requirements.txt
 pip3 install -e .
 ```
 
@@ -17,16 +20,22 @@ pip3 install -e .
 
 ### Issue Templates
 
-Please follow the [issue template on GitHub](https://github.com/inclusionAI/AReaL/tree/main/.github/ISSUE_TEMPLATE). Issues can be:
+Please follow the
+[issue template on GitHub](https://github.com/inclusionAI/AReaL/tree/main/.github/ISSUE_TEMPLATE).
+Issues can be:
+
 - Bug reports
-- Feature requests  
+- Feature requests
 - Refactor requests
 
-The required fields in the template help reduce communication overhead when resolving issues. **Issues with arbitrary formatting may be ignored.**
+The required fields in the template help reduce communication overhead when resolving
+issues. **Issues with arbitrary formatting may be ignored.**
 
 ## Pull Request Guidelines
 
-There are no specific PR templates, but **pull requests should be related to a well-templated issue**. Your PR should:
+There are no specific PR templates, but **pull requests should be related to a
+well-templated issue**. Your PR should:
+
 - Explain how the issue is resolved
 - Describe the benefits this change will provide
 - Reference the related issue number
@@ -46,24 +55,23 @@ isort . && black .
 AReaL's unit tests are based on the `pytest` framework:
 
 ```bash
-# Run all tests (excluding GPU tests)
-pytest -m "not gpu"
-
-# Run a specific test case
-pytest tests/test_something.py
+# Run all tests
+pytest -s -v areal/tests/
 ```
 
 **Note**: Running all tests may take several hours to complete.
 
 ## Documentation
 
-Writing documentation is an excellent starting point for new contributors. The documentation is located in the `docs` folder and built using [Jupyter Book](https://jupyterbook.org/en/stable/intro.html).
+Writing documentation is an excellent starting point for new contributors. The
+documentation is located in the `docs` folder and built using
+[Jupyter Book](https://jupyterbook.org/en/stable/intro.html).
 
 ### Adding New Documentation
 
 1. Create your documentation files in the `docs` folder
-2. Add the file path to `docs/_toc.yaml`
-3. Build the documentation:
+1. Add the file path to `docs/_toc.yaml`
+1. Build the documentation:
 
 ```bash
 jb build docs
@@ -71,4 +79,5 @@ jb build docs
 
 4. Preview your changes by opening the HTML files in `docs/_build/html`
 
-This process allows you to see how your documentation will appear before submitting your contribution.
\ No newline at end of file
+This process allows you to see how your documentation will appear before submitting your
+contribution.
diff --git a/docs/customization/agent.md b/docs/customization/agent.md
index 8ee9e2570..753ef81f5 100644
--- a/docs/customization/agent.md
+++ b/docs/customization/agent.md
@@ -192,7 +192,7 @@ Finally, let's complete the implementation by collecting trajectories in the
 class MultiTurnWorkflow(RolloutWorkflow):
     # ... previous methods ...
 
-    async def arun_episode(self, engine: InferenceEngine, data):
+    async def arun_episode(self, engine: InferenceEngine, data) -> TensorDict:
         # ... episode logic above ...
 
         while reward == 0 and t < self.max_turns: