From 2672d9b96bc7a941c2a35ab666190ac1b0bf19a0 Mon Sep 17 00:00:00 2001
From: hyeok9855 <hyeok9855@gmail.com>
Date: Sat, 7 Dec 2024 03:32:58 +0900
Subject: [PATCH] add docstring

---
 src/gfn/env.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/gfn/env.py b/src/gfn/env.py
index c1c1085..7249f36 100644
--- a/src/gfn/env.py
+++ b/src/gfn/env.py
@@ -94,6 +94,8 @@ def states_from_batch_shape(
 
         Args:
             batch_shape: Tuple representing the shape of the batch of states.
+            random (optional): Initalize states randomly.
+            sink (optional): States initialized with s_f (the sink state).
 
         Returns:
             States: A batch of initial states.