File size: 11,071 Bytes
c1be7c3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
ProcureRL Environment Implementation.

An OpenEnv-compliant RL environment for procurement negotiation where
an LLM agent learns to negotiate against scripted supplier opponents.
"""

import uuid
from typing import Optional, Dict, Any

try:
    from openenv.core.env_server.interfaces import Environment
except ImportError:
    Environment = object

try:
    from ..models import NegotiationAction, NegotiationObservation, NegotiationState
    from ..opponent import ScriptedPersonaOpponent
    from ..graders import grade
except ImportError:
    import sys
    import os

    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from models import NegotiationAction, NegotiationObservation, NegotiationState
    from opponent import ScriptedPersonaOpponent
    from graders import grade


TASK_CONFIG = {
    "single_issue": {
        "persona": "cooperative",
        "max_rounds": 6,
        "buyer_constraints": {
            "price": {"target": 36000, "worst": 55000, "budget": 53000}
        },
    },
    "multi_issue": {
        "persona": "cash_flow_stressed",
        "max_rounds": 8,
        "buyer_constraints": {
            "price": {"target": 40000, "worst": 58000, "budget": 55000},
            "payment_days": {"target": 60, "worst": 30, "preference": 60},
        },
    },
    "adversarial": {
        "persona": "aggressive_anchor",
        "max_rounds": 10,
        "buyer_constraints": {
            "price": {"target": 80000, "worst": 120000, "budget": 115000},
            "payment_days": {"target": 60, "worst": 30, "preference": 60},
            "support_hours": {"target": 150, "worst": 80, "preference": 150},
        },
    },
}

VALID_MOVES = ("make_offer", "accept", "reject", "bundle")


class ProcureRLEnvironment(Environment):
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        self._state = NegotiationState()
        self._opponent = None
        self._task_config = None
        self._done = False
        self._last_offer: Dict[str, Any] = {}
        self._consecutive_concessions = 0
        self._prev_agent_price: Optional[float] = None
        self._exchanges: list = []
        self._last_info: Dict[str, Any] = {}

    def reset(
        self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs
    ) -> NegotiationObservation:
        task_id = kwargs.get("task_id", "single_issue")
        seed = seed if seed is not None else 42

        if task_id not in TASK_CONFIG:
            obs = self._make_obs(
                f"Unknown task: {task_id}. Valid: {list(TASK_CONFIG.keys())}"
            )
            obs.done = True
            obs.metadata["error"] = f"unknown_task:{task_id}"
            return obs

        config = TASK_CONFIG[task_id]
        self._task_config = config
        self._done = False
        self._consecutive_concessions = 0
        self._prev_agent_price = None
        self._exchanges = []
        self._last_info = {}

        opponent_seed = hash((seed, task_id)) % (2**32)
        self._opponent = ScriptedPersonaOpponent(
            task_id=task_id, seed=opponent_seed, persona=config["persona"]
        )

        opening_msg, opening_terms = self._opponent.get_opening_message()
        self._last_offer = opening_terms
        self._opponent_opening_price = opening_terms.get("price", 52000.0)

        self._state = NegotiationState(
            task_id=task_id,
            episode_id=episode_id or str(uuid.uuid4())[:8],
            round_number=0,
            step_count=0,
            rapport_score=0.5,
            consecutive_concessions=0,
            deal_reached=False,
            final_terms=None,
            cumulative_reward=0.0,
        )

        self._exchanges.append(
            {"role": "supplier", "message": opening_msg, "terms": opening_terms}
        )

        return NegotiationObservation(
            task_id=task_id,
            round_number=0,
            max_rounds=config["max_rounds"],
            supplier_message=opening_msg,
            current_offer=opening_terms,
            last_4_exchanges=self._exchanges[-4:],
            buyer_constraints=config["buyer_constraints"],
            rapport_hint="neutral",
            done=False,
        )

    def step(self, action: NegotiationAction, **kwargs) -> NegotiationObservation:
        self._last_info = {}

        if self._done:
            obs = self._make_obs("Episode finished. Call reset().")
            obs.done = True
            obs.metadata["error"] = "episode_done"
            return obs

        if self._task_config is None:
            obs = self._make_obs("Environment not initialized. Call reset() first.")
            obs.done = True
            obs.metadata["error"] = "not_initialized"
            return obs

        if not isinstance(action, NegotiationAction):
            action_dict = (
                action if isinstance(action, dict) else {"move_type": "make_offer"}
            )
            action = NegotiationAction(
                move_type=action_dict.get("move_type", "make_offer"),
                terms=action_dict.get("terms", {}),
                message=action_dict.get("message", ""),
            )

        if action.move_type not in VALID_MOVES:
            obs = self._make_obs()
            obs.metadata["error"] = f"invalid_move_type:{action.move_type}"
            return obs

        self._state.round_number += 1
        self._state.step_count += 1
        round_num = self._state.round_number
        config = self._task_config
        max_rounds = config["max_rounds"]

        reward = 0.0

        if self._prev_agent_price is not None and "price" in action.terms:
            current_price = float(action.terms.get("price", self._prev_agent_price))
            if current_price > self._prev_agent_price:
                self._consecutive_concessions += 1
            else:
                self._consecutive_concessions = 0
        if "price" in action.terms:
            self._prev_agent_price = float(action.terms.get("price"))
        self._state.consecutive_concessions = self._consecutive_concessions

        if action.move_type in ("make_offer", "bundle"):
            opponent_msg, opponent_terms = self._opponent.respond(
                agent_message=action.message,
                agent_terms=action.terms,
                round_number=round_num,
                consecutive_concessions=self._consecutive_concessions,
            )

            self._exchanges.append(
                {"role": "agent", "message": action.message, "terms": action.terms}
            )

            if opponent_terms.get("_accepted"):
                self._done = True
                self._state.deal_reached = True
                self._state.final_terms = action.terms
                reward = grade(
                    self._state.task_id,
                    action.terms,
                    True,
                    round_num,
                    opponent_opening=self._opponent_opening_price,
                    consecutive_concessions_flag=(self._consecutive_concessions >= 2),
                )
                self._state.cumulative_reward = reward
                obs = self._make_obs(supplier_message=opponent_msg)
                obs.done = True
                obs.reward = reward
                self._last_info["deal_price"] = action.terms.get("price")
                self._exchanges.append(
                    {
                        "role": "supplier",
                        "message": opponent_msg,
                        "terms": {
                            k: v
                            for k, v in opponent_terms.items()
                            if not k.startswith("_")
                        },
                    }
                )
                return obs

            self._last_offer = {
                k: v for k, v in opponent_terms.items() if not k.startswith("_")
            }
            self._state.rapport_score = self._opponent.rapport

            self._exchanges.append(
                {"role": "supplier", "message": opponent_msg, "terms": self._last_offer}
            )

            if round_num >= max_rounds:
                self._done = True
                reward = 0.0
                obs = self._make_obs(supplier_message=opponent_msg)
                obs.done = True
                obs.reward = reward
                self._last_info["error"] = "max_rounds_reached"
                return obs

            obs = self._make_obs(supplier_message=opponent_msg)
            obs.reward = reward
            return obs

        if action.move_type == "accept":
            self._done = True
            self._state.deal_reached = True
            self._state.final_terms = self._last_offer
            reward = grade(
                self._state.task_id,
                self._last_offer,
                True,
                round_num,
                opponent_opening=self._opponent_opening_price,
                consecutive_concessions_flag=(self._consecutive_concessions >= 2),
            )
            self._state.cumulative_reward = reward
            obs = self._make_obs()
            obs.done = True
            obs.reward = reward
            self._last_info["deal_price"] = self._last_offer.get("price")
            return obs

        if action.move_type == "reject":
            if round_num >= max_rounds:
                self._done = True
                reward = 0.0
                obs = self._make_obs()
                obs.done = True
                obs.reward = reward
                self._last_info["error"] = "rejected_at_limit"
                return obs
            obs = self._make_obs()
            obs.reward = 0.0
            return obs

        obs = self._make_obs()
        obs.reward = 0.0
        return obs

    @property
    def state(self) -> NegotiationState:
        return self._state

    def close(self) -> None:
        pass

    def _make_obs(self, supplier_message: str = None) -> NegotiationObservation:
        rapport = self._state.rapport_score
        if rapport >= 0.65:
            hint = "positive"
        elif rapport <= 0.35:
            hint = "negative"
        else:
            hint = "neutral"

        return NegotiationObservation(
            task_id=self._state.task_id or "",
            round_number=self._state.round_number,
            max_rounds=self._task_config["max_rounds"] if self._task_config else 0,
            supplier_message=supplier_message or "",
            current_offer=self._last_offer,
            last_4_exchanges=self._exchanges[-4:] if self._exchanges else [],
            buyer_constraints=self._task_config["buyer_constraints"]
            if self._task_config
            else {},
            rapport_hint=hint,
            done=self._done,
            metadata=self._last_info,
        )