(fix):ReplayBuffer to pass task_name directly to add_frame method; update gym_manipulator documentation to describe environment features and usage.

This commit is contained in:
AdilZouitine
2025-06-04 14:40:29 +02:00
parent 64219571e4
commit 906dd1396d
2 changed files with 23 additions and 5 deletions

View File

@@ -547,7 +547,6 @@ class ReplayBuffer:
repo_id=repo_id,
fps=fps,
root=root,
robot=None, # TODO: (azouitine) Handle robot
robot_type=None,
features=features,
use_videos=True,
@@ -588,11 +587,8 @@ class ReplayBuffer:
else:
frame_dict[f"complementary_info.{key}"] = val
# Add task field which is required by LeRobotDataset
frame_dict["task"] = task_name
# Add to the dataset's buffer
lerobot_dataset.add_frame(frame_dict)
lerobot_dataset.add_frame(frame_dict, task=task_name)
# Move to next frame
frame_idx_in_episode += 1

View File

@@ -14,6 +14,28 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Robot Environment for LeRobot Manipulation Tasks
This module provides a comprehensive gym-compatible environment for robot manipulation
with support for:
- Multiple robot types (SO100, SO101, Koch and Moss)
- Human intervention via leader-follower control or gamepad
- End-effector and joint space control
- Image processing (cropping and resizing)
The environment is built using a composable wrapper pattern where each wrapper
adds specific functionality to the base RobotEnv.
Example:
env = make_robot_env(cfg)
obs, info = env.reset()
action = policy.select_action(obs)
obs, reward, terminated, truncated, info = env.step(action)
"""
import logging
import time
from collections import deque