Reorganize imports

cogment_lab/actors/__init__.py

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .simple import ConstantActor, RandomActor
+from cogment_lab.actors.runner import actor_runner
+from cogment_lab.actors.simple import ConstantActor, RandomActor

cogment_lab/actors/runner.py

@@ -21,7 +21,7 @@
 
 from cogment_lab.core import BaseActor
 from cogment_lab.generated import cog_settings
-from cogment_lab.utils.runners import setup_logging
+from cogment_lab.utils import setup_logging
 
 
 async def register_actor(actor: BaseActor, actor_name: str, queue: Queue, port: int = 9002):

cogment_lab/cli/__init__.py

@@ -11,3 +11,5 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from cogment_lab.cli.download_cogment import download_cogment

cogment_lab/cli/cli.py

@@ -18,7 +18,7 @@
 import logging
 import sys
 
-from cogment_lab.cli.download_cogment import download_cogment
+from cogment_lab.cli import download_cogment
 from cogment_lab.constants import COGMENT_LAB_HOME
 
 

cogment_lab/envs/__init__.py

@@ -14,3 +14,4 @@
 
 from cogment_lab.envs.gymnasium import GymEnvironment
 from cogment_lab.envs.pettingzoo import AECEnvironment, ParallelEnvironment
+from cogment_lab.envs.runner import env_runner, register_env

cogment_lab/envs/runner.py

@@ -24,7 +24,7 @@
 
 from cogment_lab.core import BaseEnv
 from cogment_lab.generated import cog_settings
-from cogment_lab.utils.runners import setup_logging
+from cogment_lab.utils import setup_logging
 
 
 async def register_env(env: BaseEnv, env_name: str, signal_queue: Queue, port: int = 9001):

cogment_lab/humans/__init__.py

@@ -11,3 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from cogment_lab.humans.actor import run_cogment_actor, start_fastapi
+from cogment_lab.humans.runner import human_actor_runner

cogment_lab/humans/runner.py

@@ -19,8 +19,8 @@
 from multiprocessing import Queue
 from typing import Any
 
-from cogment_lab.humans.actor import run_cogment_actor, start_fastapi
-from cogment_lab.utils.runners import setup_logging
+from cogment_lab.humans import run_cogment_actor, start_fastapi
+from cogment_lab.utils import setup_logging
 
 
 # def human_actor_runner(

cogment_lab/process_manager.py

@@ -29,16 +29,12 @@
 from cogment.control import Controller
 from cogment.datastore import Datastore
 
-from cogment_lab.actors.runner import actor_runner
+from cogment_lab.actors import actor_runner
 from cogment_lab.core import BaseActor, BaseEnv
-from cogment_lab.envs.runner import env_runner
+from cogment_lab.envs import env_runner
 from cogment_lab.generated import cog_settings, data_pb2
-from cogment_lab.humans.runner import human_actor_runner
-from cogment_lab.utils.trial_utils import (
-    TrialData,
-    format_data_multiagent,
-    get_actor_params,
-)
+from cogment_lab.humans import human_actor_runner
+from cogment_lab.utils import TrialData, format_data_multiagent, get_actor_params
 
 
 AgentName = str

cogment_lab/session_helpers.py

@@ -21,9 +21,13 @@
 from cogment.model_registry_v2 import ModelRegistry
 from cogment.session import ActorInfo, RecvEvent
 
-from cogment_lab.specs import AgentSpecs
-from cogment_lab.specs.action_space import Action, ActionSpace
-from cogment_lab.specs.observation_space import Observation, ObservationSpace
+from cogment_lab.specs import (
+    Action,
+    ActionSpace,
+    AgentSpecs,
+    Observation,
+    ObservationSpace,
+)
 
 
 class ActorSessionHelper:

cogment_lab/specs/__init__.py

@@ -12,5 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .encode_rendered_frame import encode_rendered_frame
-from .environment_specs import AgentSpecs
+from cogment_lab.specs.action_space import Action, ActionSpace
+from cogment_lab.specs.encode_rendered_frame import (
+    decode_rendered_frame,
+    encode_rendered_frame,
+)
+from cogment_lab.specs.environment_specs import AgentSpecs
+from cogment_lab.specs.ndarray_serialization import (
+    SerializationFormat,
+    deserialize_ndarray,
+    serialize_ndarray,
+)
+from cogment_lab.specs.observation_space import Observation, ObservationSpace
+from cogment_lab.specs.spaces_serialization import (
+    deserialize_space,
+    serialize_gym_space,
+)

cogment_lab/specs/action_space.py

@@ -17,8 +17,7 @@
 import gymnasium as gym
 
 from cogment_lab.generated.data_pb2 import PlayerAction  # type: ignore
-
-from .ndarray_serialization import deserialize_ndarray, serialize_ndarray
+from cogment_lab.specs import deserialize_ndarray, serialize_ndarray
 
 
 # pylint: disable=attribute-defined-outside-init

cogment_lab/specs/environment_specs.py

@@ -16,13 +16,15 @@
 
 import gymnasium as gym
 
+from cogment_lab.constants import DEFAULT_RENDERED_WIDTH
 from cogment_lab.generated.data_pb2 import AgentSpecs as PbAgentSpecs  # type: ignore
-
-from ..constants import DEFAULT_RENDERED_WIDTH
-from .action_space import ActionSpace
-from .ndarray_serialization import SerializationFormat
-from .observation_space import ObservationSpace
-from .spaces_serialization import deserialize_space, serialize_gym_space
+from cogment_lab.specs import (
+    ActionSpace,
+    ObservationSpace,
+    SerializationFormat,
+    deserialize_space,
+    serialize_gym_space,
+)
 
 
 class AgentSpecs:

cogment_lab/specs/observation_space.py

@@ -17,9 +17,12 @@
 
 from cogment_lab.constants import DEFAULT_RENDERED_WIDTH
 from cogment_lab.generated.data_pb2 import Observation as PbObservation  # type: ignore
-
-from .encode_rendered_frame import decode_rendered_frame, encode_rendered_frame
-from .ndarray_serialization import deserialize_ndarray, serialize_ndarray
+from cogment_lab.specs import (
+    decode_rendered_frame,
+    deserialize_ndarray,
+    encode_rendered_frame,
+    serialize_ndarray,
+)
 
 
 # pylint: disable=attribute-defined-outside-init

cogment_lab/specs/spaces_serialization.py

@@ -21,8 +21,7 @@
 from cogment_lab.generated.spaces_pb2 import MultiBinary  # type: ignore
 from cogment_lab.generated.spaces_pb2 import MultiDiscrete  # type: ignore
 from cogment_lab.generated.spaces_pb2 import Space  # type: ignore
-
-from .ndarray_serialization import (
+from cogment_lab.specs import (
     SerializationFormat,
     deserialize_ndarray,
     serialize_ndarray,

cogment_lab/utils/__init__.py

@@ -12,4 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .import_class import import_object
+from cogment_lab.utils.grpc import extend_actor_config
+from cogment_lab.utils.import_class import import_object
+from cogment_lab.utils.runners import process_cleanup, setup_logging
+from cogment_lab.utils.trial_utils import (
+    TrialData,
+    format_data_multiagent,
+    get_actor_params,
+)

cogment_lab/utils/coltra_utils.py

@@ -17,7 +17,7 @@
 from coltra import Agent
 from coltra.buffers import Action, Observation, OnPolicyRecord
 
-from cogment_lab.utils.trial_utils import TrialData
+from cogment_lab.utils import TrialData
 
 
 def convert_trial_data_to_coltra(trial_data: TrialData, agent: Agent) -> OnPolicyRecord:

cogment_lab/utils/trial_utils.py

@@ -27,7 +27,7 @@
 
 from cogment_lab.generated import cog_settings
 from cogment_lab.specs import AgentSpecs
-from cogment_lab.utils.grpc import extend_actor_config
+from cogment_lab.utils import extend_actor_config
 
 
 def get_actor_params(

examples/docs_example.py

@@ -18,7 +18,7 @@
 from cogment_lab import Cogment
 from cogment_lab.actors import ConstantActor, RandomActor
 from cogment_lab.envs import GymEnvironment
-from cogment_lab.utils.runners import process_cleanup
+from cogment_lab.utils import process_cleanup
 
 
 LUNAR_LANDER_ACTIONS = ["no-op", "ArrowRight", "ArrowUp", "ArrowLeft"]

examples/gymnasium/interactive.py

@@ -18,7 +18,7 @@
 from cogment_lab.actors import ConstantActor, RandomActor
 from cogment_lab.envs.gymnasium import GymEnvironment
 from cogment_lab.process_manager import Cogment
-from cogment_lab.utils.trial_utils import format_data_multiagent
+from cogment_lab.utils import format_data_multiagent
 
 
 async def main():