File size: 4,371 Bytes
aab0192
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
"""
client.py -- Typed Python client for HypothesisLab.

Built on openenv.core.env_client.EnvClient (WebSocket-based, persistent).
This is what the RL trainer and baseline inference script import.
"""

from __future__ import annotations

from typing import Any, Dict, Optional

try:
    from openenv.core.env_client import EnvClient
    from openenv.core.client_types import StepResult
except ImportError:
    raise ImportError(
        "openenv-core is required. Install with: pip install openenv-core"
    )

try:
    from .models import (
        ActionType,
        ExperimentType,
        HypLabAction,
        HypLabObservation,
        HypLabState,
        NoiseLevelTag,
    )
except ImportError:
    from models import (
        ActionType,
        ExperimentType,
        HypLabAction,
        HypLabObservation,
        HypLabState,
        NoiseLevelTag,
    )


class HypothesisLabEnv(EnvClient[HypLabAction, HypLabObservation, HypLabState]):
    """
    Typed async client for the Scientific Hypothesis Lab environment.

    Usage (async):
        async with HypothesisLabEnv(base_url="http://localhost:8000") as env:
            result = await env.reset(noise_level="low", domain="physics")
            obs = result.observation
            ...

    Usage (sync):
        env = HypothesisLabEnv(base_url="http://localhost:8000").sync()
        with env:
            result = env.reset(noise_level="low")
            ...
    """

    def _step_payload(self, action: HypLabAction) -> Dict[str, Any]:
        return action.model_dump(exclude_none=True)

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[HypLabObservation]:
        obs_data = payload.get("observation", payload)
        obs = HypLabObservation(**obs_data)
        return StepResult(
            observation=obs,
            reward=payload.get("reward", obs.reward or 0.0),
            done=payload.get("done", obs.done),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> HypLabState:
        return HypLabState(**payload)

    async def run_intervention(
        self,
        control_variable: str,
        control_value: float,
        target_variable: str,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.INTERVENTION,
            control_variable=control_variable,
            target_variable=target_variable,
            control_value=control_value,
        )
        return await self.step(action)

    async def run_correlation(
        self,
        control_variable: str,
        control_range: list[float],
        target_variable: str,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.CORRELATION,
            control_variable=control_variable,
            control_range=control_range,
            target_variable=target_variable,
        )
        return await self.step(action)

    async def run_counterfactual(
        self,
        control_variable: str,
        delta: float,
        target_variable: str,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.COUNTERFACTUAL,
            control_variable=control_variable,
            control_value=delta,
            target_variable=target_variable,
        )
        return await self.step(action)

    async def run_passive(
        self, target_variable: str
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.PASSIVE,
            target_variable=target_variable,
            control_variable=target_variable,
        )
        return await self.step(action)

    async def submit_hypothesis(
        self,
        hypothesis_text: str,
        hypothesis_equations: Optional[list[str]] = None,
        confidence: float = 0.75,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.SUBMIT,
            hypothesis_text=hypothesis_text,
            hypothesis_equations=hypothesis_equations,
            confidence=max(0.0, min(1.0, confidence)),
        )
        return await self.step(action)