diff --git a/src/aviary/core.py b/src/aviary/core.py
index d28be1ae..12f400e9 100644
--- a/src/aviary/core.py
+++ b/src/aviary/core.py
@@ -9,7 +9,7 @@
     TaskDataset,
 )
 from aviary.env_client import EnvironmentClient
-from aviary.message import MalformedMessageError, Message, join
+from aviary.message import EnvStateMessage, MalformedMessageError, Message, join
 from aviary.render import Renderer
 from aviary.tools import (
     INVALID_TOOL_NAME,
@@ -38,6 +38,7 @@
     "DummyEnv",
     "DummyEnvState",
     "DummyTaskDataset",
+    "EnvStateMessage",
     "Environment",
     "EnvironmentClient",
     "Frame",