Polish documents

Mossforest · Jun 21, 2024 · c901c2b · c901c2b
1 parent 4c7ff2b
commit c901c2b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/docs/source/api_doc/neural_network/index.rst b/docs/source/api_doc/neural_network/index.rst
@@ -37,7 +37,7 @@ DiT
 
 DiT1D
 ------
-.. autoclass:: DiT2D
+.. autoclass:: DiT1D
     :special-members: __init__
     :members:
 

diff --git a/grl/agents/__init__.py b/grl/agents/__init__.py
@@ -38,4 +38,5 @@ def action_transform(action, return_as_torch_tensor: bool = False):
 
 from .base import BaseAgent
 from .qgpo import QGPOAgent
+from .srpo import SRPOAgent
 from .gm import GPAgent
diff --git a/grl/neural_network/__init__.py b/grl/neural_network/__init__.py
@@ -495,7 +495,7 @@ def forward(
         return result
 
 
-from .transformers.dit import DiT, DiT1D, DiT3D
+from .transformers.dit import DiT, DiT1D, DiT2D, DiT3D 
 
 
 class Sequential(nn.Module):