pushing brave search multi turn rl code

[pranavraja99]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces a new BrowseEnv for multi-turn reinforcement learning using Brave Search, along with the necessary tool integrations and an example script. The overall structure is good, but there are several critical issues that need addressing. The BraveSearchToolGroup implementation incorrectly handles API requests by not using the configured session or URL, and its retry logic is flawed. The new BrowseEnv has a bug in its step method that can cause crashes when no tool is called. Additionally, there are some unused parameters and leftover debugging print statements that should be cleaned up. I've provided specific comments and suggestions to resolve these issues.

[gemini-code-assist]

The _search method has several critical issues that deviate from the class design and introduce bugs:

1. Session Not Used: It does not use the self.session object for making requests, instead calling requests.get() directly. This negates the benefits of the connection pooling configured in __init__.
2. URL Hardcoded: It uses a hardcoded URL to the Brave Search API and ignores self.search_url from __init__. This makes the search_url parameter misleading, especially since the example scripts configure it for a local retriever.

This method should be refactored to use self.session.get(self.search_url, ...) to align with the class's design and configuration.

[gemini-code-assist]

The compute_score function is being imported from skyrl_gym.envs.search.utils, but a new utils.py file with this function is also being added in the browse directory. To ensure the correct utility is used and to follow better module structure, you should import from the local utils.py file.

Suggested change

	from skyrl_gym.envs.search.utils import compute_score
	from .utils import compute_score

[gemini-code-assist]

The current implementation of the step method has a bug where it will raise an exception if the model's action does not contain a <search> tag. This is because _parse_action returns [None], and _execute_tool is then called with None as input, which is not handled. The agent should be able to produce a response without calling a tool.

The logic should be updated to only execute a tool if a valid tool call is parsed from the action. I've also refactored to initialize info only when a tool is actually called.

        query = self._parse_action(action)
        observation = None
        info = {}

        if query[0] is not None:
            try:
                observation = self._execute_tool("BraveSearchToolGroup", "browse", query)
                info = {
                    "tool_group": "BraveSearchToolGroup",
                    "tool_name": "browse",
                    "tool_input": query,
                }
            except Exception as e:
                error = str(e)
                observation = None

        # Wrap the observation properly as a message
        if observation:
            new_obs = {"role": "user", "content": observation}
        elif error:
            # Give error as observation if any
            new_obs = {"role": "user", "content": error}
        else:
            new_obs = None

        # Update chat history
        if new_obs:
            self.chat_history.append(new_obs)

        return BaseTextEnvStepOutput(
            observations=[new_obs] if new_obs else [],
            reward=reward,
            done=done,
            metadata=info,
        )

[gemini-code-assist]

The retry logic in _search has several issues:

1. Infinite Loop: The while True loop for retries is not bounded, which could lead to an infinite loop on repeated 429 errors. The MAX_RETRIES constant should be used to limit retries.
2. Hardcoded Timeout: The request timeout is hardcoded to 30 instead of using self.timeout.
3. Logging: print is used for logging rate limit warnings. It's better to use the configured logger (e.g., logger.warning(...)).
4. Unused Constants: MAX_RETRIES and INITIAL_RETRY_DELAY are defined but not used.

The retry logic should be updated to use a bounded loop and use the class attributes for configuration.

[gemini-code-assist]

The method parameter is defined but not used within this function. It should be removed to avoid confusion.

Suggested change

	def compute_score(solution_str, ground_truth, method="strict", format_score=0.0, score=1.0):
	def compute_score(solution_str, ground_truth, format_score=0.0, score=1.0):

[gemini-code-assist]

The method parameter is defined but not used within this function. It should be removed to avoid confusion.

Suggested change

	def compute_score_subem(solution_str, ground_truth, method="strict", format_score=0.0, score=1.0):
	def compute_score_subem(solution_str, ground_truth, format_score=0.0, score=1.0):

[gemini-code-assist]

There's a mismatch between the run_name which specifies '8turns' and the generator.max_turns parameter which is set to 6 on line 47. To avoid confusion when tracking experiments, these should be consistent.

Suggested change

	trainer.run_name="skyrl-browse_8turns_maxgeneratelen_500" \
	trainer.run_name="skyrl-browse_6turns_maxgeneratelen_500" \

[gemini-code-assist]

These print statements appear to be for debugging and should be removed before merging.

Suggested change

	print('self.cfg.environment.env_class',self.cfg.environment.env_class)
	generator: GeneratorInterface = self.get_generator(self.cfg, tokenizer, inference_engine_client)
	print('GOT TO GENERATOR')
	generator: GeneratorInterface = self.get_generator(self.cfg, tokenizer, inference_engine_client)